diff --git a/api/services/agent/chat/base_runner.py b/api/services/agent/chat/base_runner.py index 719d4db7419946..a21a5e9546efbc 100644 --- a/api/services/agent/chat/base_runner.py +++ b/api/services/agent/chat/base_runner.py @@ -412,9 +412,9 @@ def organize_agent_history( else: tool_responses = dict.fromkeys(tool_names, observation_payload) - for tool in tool_names: - # generate a uuid for tool call - tool_call_id = str(uuid.uuid4()) + for index, tool in enumerate(tool_names): + # Keep IDs stable across turns so replayed history can use prompt caching. + tool_call_id = str(uuid.uuid5(uuid.NAMESPACE_OID, f"{agent_thought.id}:{index}:{tool}")) tool_calls.append( AssistantPromptMessage.ToolCall( id=tool_call_id, diff --git a/api/services/workflow/execution/adapters/response_converter.py b/api/services/workflow/execution/adapters/response_converter.py index eaf506cffa997c..3ceca89a6ad6b7 100644 --- a/api/services/workflow/execution/adapters/response_converter.py +++ b/api/services/workflow/execution/adapters/response_converter.py @@ -275,6 +275,7 @@ def workflow_start_to_stream_response( run_id = self._ensure_workflow_run_id(workflow_run_id) started_at = naive_utc_now() self._workflow_started_at = started_at + self._application_generate_entity.workflow_execution_timing.start() if reason == WorkflowStartReason.RESUMPTION: self._node_snapshots.update( (NodeExecutionId(snapshot.execution_id), snapshot) for snapshot in node_execution_snapshots @@ -310,7 +311,7 @@ def workflow_finish_to_stream_response( ) finished_at = naive_utc_now() - elapsed_time = (finished_at - started_at).total_seconds() + elapsed_time = self._application_generate_entity.workflow_execution_timing.finish() outputs_mapping = graph_runtime_state.outputs or {} encoded_outputs = WorkflowRuntimeTypeConverter().to_json_encodable(outputs_mapping) @@ -363,8 +364,7 @@ def workflow_pause_to_stream_response( raise ValueError( "workflow_pause_to_stream_response called before workflow_start_to_stream_response", ) - paused_at = naive_utc_now() - elapsed_time = (paused_at - started_at).total_seconds() + elapsed_time = self._application_generate_entity.workflow_execution_timing.finish() encoded_outputs = self._encode_outputs(event.outputs) or {} if self._application_generate_entity.invoke_from == InvokeFrom.SERVICE_API: encoded_outputs = {} diff --git a/api/services/workflow/execution/adapters/workflow/app_generator.py b/api/services/workflow/execution/adapters/workflow/app_generator.py index 71a0043fbc97b9..14e0a92fadf3b6 100644 --- a/api/services/workflow/execution/adapters/workflow/app_generator.py +++ b/api/services/workflow/execution/adapters/workflow/app_generator.py @@ -42,7 +42,6 @@ from models.enums import WorkflowRunTriggeredFrom from models.model import App, EndUser from models.workflow import WorkflowNodeExecutionTriggeredFrom -from repositories.entities.workflow_execution import WorkflowExecutionIdentity from services.app.file_input_defaults import grant_file_default_access from services.app.generation.input_adapter import AppInputAdapter from services.workflow.contracts import WorkflowSnapshot @@ -60,6 +59,7 @@ if TYPE_CHECKING: from services.workflow.execution.contracts import WorkflowRunDriver + from services.workflow_execution_identity import WorkflowExecutionIdentity logger = logging.getLogger(__name__) diff --git a/api/tests/unit_tests/core/agent/conftest.py b/api/tests/unit_tests/core/agent/conftest.py index a2aa50172093f6..34aec09d569209 100644 --- a/api/tests/unit_tests/core/agent/conftest.py +++ b/api/tests/unit_tests/core/agent/conftest.py @@ -1,5 +1,9 @@ +from unittest.mock import MagicMock, create_autospec + import pytest +from services.agent.chat.ports import AgentToolInvoker + class DummyTool: def __init__(self, name): @@ -78,3 +82,8 @@ def _factory(**kwargs): return DummyScratchpadUnit(**kwargs) return _factory + + +@pytest.fixture +def agent_tool_invoker() -> MagicMock: + return create_autospec(AgentToolInvoker, instance=True, spec_set=True) diff --git a/api/tests/unit_tests/core/agent/test_base_agent_runner.py b/api/tests/unit_tests/core/agent/test_base_agent_runner.py index 6f3c23e9d329a6..dc8fab1fdec28b 100644 --- a/api/tests/unit_tests/core/agent/test_base_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_base_agent_runner.py @@ -1,13 +1,13 @@ import json from datetime import datetime from decimal import Decimal +from unittest.mock import MagicMock, create_autospec import pytest from pytest_mock import MockerFixture -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker -import core.agent.base_agent_runner as module -from core.agent.base_agent_runner import BaseAgentRunner +import services.agent.chat.base_runner as module from core.agent.entities import AgentEntity, AgentToolEntity from core.app.app_config.entities import ( AppAdditionalFeatures, @@ -34,16 +34,28 @@ ToolParameter, ToolProviderType, ) -from core.tools.utils.dataset_retriever.dataset_multi_retriever_tool import DatasetMultiRetrieverTool -from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool +from extensions.application_services.workflow import WorkflowExecutionDependencies from extensions.ext_storage import storage from extensions.storage.storage_type import StorageType from graphon.file import FileTransferMethod, FileType from graphon.model_runtime.entities import LLMUsage, PromptMessageTool from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo -from models.model import AppMode, AppModelConfig, Conversation, Message, MessageAgentThought, MessageFile, UploadFile -from services.workflow_run_agg import WorkflowRunAgg +from models.model import ( + AppMode, + AppModelConfig, + Conversation, + Message, + MessageAgentThought, + MessageFile, + UploadFile, +) +from repositories.app.generation_repository import AppGenerationRepository +from services.agent.chat.base_runner import BaseAgentRunner +from services.app.generation.ports import AgentHistoryMessage +from services.knowledge.retrieval.dataset_retrieval import DatasetRetrieval +from services.tools.dataset.tool import DatasetRetrieverTool from tests.helpers.unit.core.model_fixtures import make_model_config +from tests.helpers.unit.model_factories import make_app from tests.tool_fixtures import make_runtime_tool @@ -112,7 +124,10 @@ def _thought( def _persist_history(session: Session, *messages: Message, thoughts: list[MessageAgentThought] | None = None) -> None: - session.add_all([_conversation(), *messages, *(thoughts or [])]) + for message in messages: + message.created_at = datetime(2030, 1, 1) + session.merge(message) + session.add_all(thoughts or []) session.commit() @@ -190,26 +205,24 @@ def _tool_entity(name: str) -> ToolEntity: ) -def _dataset_tool(name: str) -> DatasetRetrieverTool: +def _dataset_tool(mocker: MockerFixture, name: str, _sessions) -> DatasetRetrieverTool: return DatasetRetrieverTool( entity=_tool_entity(name), runtime=ToolRuntime(tenant_id="tenant"), - retrieval_tool=DatasetMultiRetrieverTool( - tenant_id="tenant", - dataset_ids=[], - return_resource=False, - retriever_from="dev", - reranking_provider_name="", - reranking_model_name="", - ), + retrieval=create_autospec(DatasetRetrieval, instance=True, spec_set=True), + dataset_id="dataset", + config=DatasetRetrieveConfigEntity(retrieve_strategy="single"), + top_k=2, + inputs={}, + invoke_from=InvokeFrom.DEBUGGER, + return_resource=False, + hit_callback=mocker.Mock(), + app_id="app1", ) @pytest.fixture -def database_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Bind legacy global-session writes to a real SQLite Session.""" - - monkeypatch.setattr(module.db, "session", sqlite_session) +def database_session(sqlite_session: Session) -> Session: return sqlite_session @@ -228,7 +241,13 @@ def queue_manager(mocker: MockerFixture) -> MessageBasedAppQueueManager: @pytest.fixture def runner( - sqlite_session: Session, mocker: MockerFixture, queue_manager: MessageBasedAppQueueManager + sqlite_session: Session, + mocker: MockerFixture, + workflow_runtime: WorkflowExecutionDependencies, + queue_manager: MessageBasedAppQueueManager, + agent_tool_invoker: MagicMock, + *, + app_records: AppGenerationRepository, ) -> BaseAgentRunner: app_config = _app_config() model_config = make_model_config(provider="provider", model="model", mode="chat") @@ -239,9 +258,14 @@ def runner( ) mocker.patch.object(model_instance.model_type_instance, "get_model_schema", return_value=None) + sqlite_session.add_all([make_app(app_id="app1", tenant_id="tenant"), _conversation(), _message()]) + sqlite_session.commit() + return BaseAgentRunner( - execution_driver=WorkflowRunAgg.run, - session=sqlite_session, + dataset_tools=workflow_runtime.dataset_tools, + records=app_records, + workflow_runtime=workflow_runtime, + tool_invoker=agent_tool_invoker, tenant_id="tenant", application_generate_entity=_app_generate(app_config=app_config), conversation=_conversation(), @@ -308,6 +332,8 @@ def test_create_agent_thought_persists_row( files: list[str], expected_files: str, ) -> None: + database_session.add(_message(message_id="message-1")) + database_session.commit() thought_id = runner.create_agent_thought("message-1", "message", "tool", "input", files) stored = database_session.get(MessageAgentThought, thought_id) @@ -319,7 +345,7 @@ def test_create_agent_thought_persists_row( def _persist_thought(session: Session, *, tool: str = "tool1;tool2") -> MessageAgentThought: - thought = _thought(thought_id="thought-1", tool=tool, thought="") + thought = _thought(thought_id="thought-1", message_id="msg_current", tool=tool, thought="") session.add(thought) session.commit() return thought @@ -366,6 +392,7 @@ def test_save_agent_thought_full_update( usage, ) + database_session.expire_all() stored = database_session.get(MessageAgentThought, thought.id) assert stored is not None assert stored.answer == "answer" @@ -386,6 +413,7 @@ def test_save_agent_thought_uses_label_fallback( runner.save_agent_thought(thought.id, None, None, None, None, None, None, [], None) + database_session.expire_all() stored = database_session.get(MessageAgentThought, thought.id) assert stored is not None assert json.loads(stored.tool_labels_str)["unknown_tool"]["en_US"] == "unknown_tool" @@ -401,6 +429,7 @@ def test_save_agent_thought_preserves_existing_labels( runner.save_agent_thought(thought.id, None, None, None, None, None, None, [], None) + database_session.expire_all() stored = database_session.get(MessageAgentThought, thought.id) assert stored is not None assert json.loads(stored.tool_labels_str)["tool1"]["en_US"] == "existing" @@ -418,10 +447,10 @@ def test_save_agent_thought_serialization_fallbacks( tool_meta = {"c": 3} real_dumps = json.dumps - def dumps_side_effect(value, *args, **kwargs): - if value in (tool_input, observation, tool_meta) and kwargs.get("ensure_ascii") is False: + def dumps_side_effect(value: object, *, ensure_ascii: bool = True) -> str: + if value in (tool_input, observation, tool_meta) and not ensure_ascii: raise TypeError("fail") - return real_dumps(value, *args, **kwargs) + return real_dumps(value, ensure_ascii=ensure_ascii) mocker.patch.object(module.json, "dumps", side_effect=dumps_side_effect) @@ -437,6 +466,7 @@ def dumps_side_effect(value, *args, **kwargs): None, ) + database_session.expire_all() stored = database_session.get(MessageAgentThought, thought.id) assert stored is not None assert isinstance(stored.tool_input, str) @@ -454,6 +484,7 @@ def test_save_agent_thought_accepts_empty_message_ids( thought = _persist_thought(database_session) runner.save_agent_thought(thought.id, None, None, None, None, "meta_string", None, messages_ids, None) # type: ignore[arg-type] + database_session.expire_all() stored = database_session.get(MessageAgentThought, thought.id) assert stored is not None assert stored.tool_meta_str == "meta_string" @@ -508,8 +539,24 @@ def _upload_file() -> UploadFile: ) +def _history_entry(session: Session) -> AgentHistoryMessage: + from sqlalchemy import select + + from models.model import load_annotation_reply_config + + message = _message(message_id="m1") + files = list(session.scalars(select(MessageFile).where(MessageFile.message_id == message.id))) + config = message.app_model_config_with_session(session=session) + return AgentHistoryMessage( + message=message, + thoughts=[], + files=files, + model_config=config.to_dict(annotation_reply=load_annotation_reply_config(session, "app1")) if config else None, + ) + + def test_organize_user_prompt_without_files(runner: BaseAgentRunner, sqlite_session: Session) -> None: - result = runner.organize_agent_user_prompt(_message(message_id="m1"), session=sqlite_session) + result = runner.organize_agent_user_prompt(_history_entry(sqlite_session)) assert result.content == "hello" @@ -517,7 +564,7 @@ def test_organize_user_prompt_with_files_but_no_config(runner: BaseAgentRunner, sqlite_session.add(_message_file()) sqlite_session.commit() - result = runner.organize_agent_user_prompt(_message(message_id="m1"), session=sqlite_session) + result = runner.organize_agent_user_prompt(_history_entry(sqlite_session)) assert result.content == "hello" @@ -527,12 +574,12 @@ def test_organize_user_prompt_uses_file_config( sqlite_session: Session, ) -> None: config = _file_enabled_app_model_config() - sqlite_session.add_all([config, _conversation(app_model_config_id=config.id), _message_file()]) + sqlite_session.merge(_conversation(app_model_config_id=config.id)) + sqlite_session.add_all([config, _message_file()]) sqlite_session.commit() result = runner.organize_agent_user_prompt( - _message(message_id="m1"), - session=sqlite_session, + _history_entry(sqlite_session), ) assert result.content == "hello" @@ -549,13 +596,13 @@ def test_organize_user_prompt_builds_file_content( belongs_to=MessageFileBelongsTo.USER, upload_file_id=upload_file.id, ) - sqlite_session.add_all([config, _conversation(app_model_config_id=config.id), upload_file, message_file]) + sqlite_session.merge(_conversation(app_model_config_id=config.id)) + sqlite_session.add_all([config, upload_file, message_file]) sqlite_session.commit() mocker.patch.object(storage, "load", return_value=b"image") result = runner.organize_agent_user_prompt( - _message(message_id="m1"), - session=sqlite_session, + _history_entry(sqlite_session), ) assert isinstance(result.content, list) @@ -563,21 +610,21 @@ def test_organize_user_prompt_builds_file_content( assert isinstance(result.content[-1], module.TextPromptMessageContent) -def test_organize_history_empty_preserves_system_prompt(runner: BaseAgentRunner, sqlite_session: Session) -> None: +def test_organize_history_empty_preserves_system_prompt(runner: BaseAgentRunner) -> None: system_message = module.SystemPromptMessage(content="sys") - result = runner.organize_agent_history([system_message], session=sqlite_session) + result = runner.organize_agent_history([system_message], history=[]) assert result == [system_message] def test_organize_history_with_answer_only(runner: BaseAgentRunner, sqlite_session: Session) -> None: _persist_history(sqlite_session, _message(message_id="m1", answer="answer")) - result = runner.organize_agent_history([], session=sqlite_session) + result = runner.organize_agent_history([], history=runner._records.agent_history(runner._identity)) assert any(isinstance(item, module.AssistantPromptMessage) and item.content == "answer" for item in result) def test_organize_history_skips_current_message(runner: BaseAgentRunner, sqlite_session: Session) -> None: _persist_history(sqlite_session, _message(message_id="msg_current", answer="answer")) - result = runner.organize_agent_history([], session=sqlite_session) + result = runner.organize_agent_history([], history=runner._records.agent_history(runner._identity)) assert result == [] @@ -607,7 +654,7 @@ def test_organize_history_reconstructs_tool_flows( ) _persist_history(sqlite_session, message, thoughts=[thought]) - result = runner.organize_agent_history([], session=sqlite_session) + result = runner.organize_agent_history([], history=runner._records.agent_history(runner._identity)) assert isinstance(result, list) assert any(isinstance(item, module.AssistantPromptMessage) for item in result) @@ -627,8 +674,9 @@ def test_organize_history_reuses_stable_tool_call_ids( ) _persist_history(sqlite_session, message, thoughts=[thought]) - first = runner.organize_agent_history([], session=sqlite_session) - second = runner.organize_agent_history([], session=sqlite_session) + history = runner._records.agent_history(runner._identity) + first = runner.organize_agent_history([], history=history) + second = runner.organize_agent_history([], history=history) def tool_call_ids(messages: list[module.PromptMessage]) -> tuple[list[str], list[str]]: call_ids = [ @@ -653,7 +701,7 @@ def test_organize_history_without_tool_name(runner: BaseAgentRunner, sqlite_sess thought = _thought(thought_id="thought-1", message_id=message.id, tool=None) _persist_history(sqlite_session, message, thoughts=[thought]) - result = runner.organize_agent_history([], session=sqlite_session) + result = runner.organize_agent_history([], history=runner._records.agent_history(runner._identity)) assert any(isinstance(item, module.AssistantPromptMessage) and item.content == "thinking" for item in result) @@ -674,8 +722,10 @@ def test_convert_tool_to_prompt_message_tool(runner: BaseAgentRunner, mocker: Mo assert prompt_tool.parameters == schema -def test_convert_dataset_retriever_tool(runner: BaseAgentRunner) -> None: - dataset_tool = _dataset_tool("ds") +def test_convert_dataset_retriever_tool( + runner: BaseAgentRunner, mocker: MockerFixture, sqlite_session_factory: sessionmaker[Session] +) -> None: + dataset_tool = _dataset_tool(mocker, "ds", sqlite_session_factory) prompt = runner._convert_dataset_retriever_tool_to_prompt_message_tool(dataset_tool) @@ -683,11 +733,13 @@ def test_convert_dataset_retriever_tool(runner: BaseAgentRunner) -> None: assert prompt.parameters["required"] == ["query"] -def test_init_prompt_tools_adds_agent_and_dataset_tools(runner: BaseAgentRunner, mocker: MockerFixture) -> None: +def test_init_prompt_tools_adds_agent_and_dataset_tools( + runner: BaseAgentRunner, mocker: MockerFixture, sqlite_session_factory: sessionmaker[Session] +) -> None: agent_tool = _agent_tool("agent_tool") agent_runtime = _tool("agent_tool", []) mocker.patch.object(module.ToolManager, "get_agent_tool_runtime", return_value=agent_runtime) - dataset_tool = _dataset_tool("dataset_tool") + dataset_tool = _dataset_tool(mocker, "dataset_tool", sqlite_session_factory) runner.app_config.agent = _agent(agent_tool) runner.dataset_tools = [dataset_tool] @@ -712,6 +764,10 @@ def test_init_uses_real_session_for_count_and_dependencies( sqlite_session: Session, mocker: MockerFixture, queue_manager: MessageBasedAppQueueManager, + agent_tool_invoker: MagicMock, + workflow_runtime: WorkflowExecutionDependencies, + *, + app_records: AppGenerationRepository, ) -> None: sqlite_session.add_all( [ @@ -721,11 +777,7 @@ def test_init_uses_real_session_for_count_and_dependencies( ] ) sqlite_session.commit() - get_dataset_tools = mocker.patch.object( - module.DatasetRetrieverTool, - "get_dataset_tools", - return_value=["ds_tool"], - ) + get_dataset_tools = mocker.Mock(return_value=["ds_tool"]) model_config = make_model_config(provider="provider", model="m", mode="chat") model_config.model_schema.features = [module.ModelFeature.STREAM_TOOL_CALL, module.ModelFeature.VISION] model_instance = ModelInstance( @@ -746,9 +798,13 @@ def test_init_uses_real_session_for_count_and_dependencies( app_generate = _app_generate(app_config=app_config, files=["file1"]) message = _message(message_id="msg1") + sqlite_session.add_all([make_app(app_id="app1", tenant_id="tenant"), _conversation(), message]) + sqlite_session.commit() initialized = BaseAgentRunner( - execution_driver=WorkflowRunAgg.run, - session=sqlite_session, + dataset_tools=get_dataset_tools, + records=app_records, + tool_invoker=agent_tool_invoker, + workflow_runtime=workflow_runtime, tenant_id="tenant", application_generate_entity=app_generate, conversation=_conversation(), @@ -766,4 +822,6 @@ def test_init_uses_real_session_for_count_and_dependencies( assert initialized.dataset_tools == ["ds_tool"] assert initialized.agent_thought_count == 2 assert initialized.history_prompt_messages == [] - assert get_dataset_tools.call_args.kwargs["session"] is sqlite_session + assert initialized._workflow_runtime is workflow_runtime + assert "session" not in get_dataset_tools.call_args.kwargs + assert get_dataset_tools.call_args.kwargs["tenant_id"] == "tenant" diff --git a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py index 635ddcfa803eb6..cae6c50885df72 100644 --- a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py @@ -9,7 +9,6 @@ from sqlalchemy.engine import Engine from sqlalchemy.orm import Session -from core.agent.cot_agent_runner import CotAgentRunner from core.agent.entities import AgentScratchpadUnit from core.agent.errors import AgentMaxIterationError from core.app.entities.app_invoke_entities import CreditUsageCreatedBy @@ -18,6 +17,7 @@ from libs.datetime_utils import naive_utc_now from models.enums import ConversationFromSource, MessageStatus from models.model import AppMode, Conversation, Message +from services.agent.chat.cot_runner import CotAgentRunner def _make_conversation(*, conversation_id: str = "conv1") -> Conversation: @@ -66,10 +66,12 @@ def _organize_prompt_messages(self): @pytest.fixture -def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[DummyRunner]: +def runner( + mocker: MockerFixture, sqlite_engine: Engine, app_records, agent_tool_invoker: MagicMock +) -> Iterator[DummyRunner]: # Prevent BaseAgentRunner __init__ from hitting database mocker.patch( - "core.agent.base_agent_runner.BaseAgentRunner.organize_agent_history", + "services.agent.chat.base_runner.BaseAgentRunner.organize_agent_history", return_value=[], ) # Prepare required constructor dependencies for BaseAgentRunner @@ -99,6 +101,7 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[DummyRunner message = _make_message() runner = DummyRunner( + _tool_invoker=agent_tool_invoker, tenant_id="tenant", application_generate_entity=application_generate_entity, conversation=_make_conversation(), @@ -124,6 +127,7 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[DummyRunner runner.session = Session(sqlite_engine) try: + runner._records = app_records yield runner finally: runner.session.close() @@ -145,18 +149,6 @@ def test_fill_in_inputs(self, runner: DummyRunner, instruction, inputs, expected assert result == expected -class TestConvertDictToAction: - def test_convert_valid_dict(self, runner: DummyRunner): - action_dict = {"action": "test", "action_input": {"a": 1}} - action = runner._convert_dict_to_action(action_dict) - assert action.action_name == "test" - assert action.action_input == {"a": 1} - - def test_convert_missing_keys(self, runner: DummyRunner): - with pytest.raises(KeyError): - runner._convert_dict_to_action({"invalid": 1}) - - class TestFormatAssistantMessage: def test_format_assistant_message_multiple_scratchpads(self, runner: DummyRunner): sp1 = AgentScratchpadUnit( @@ -208,27 +200,24 @@ def test_format_with_action_and_observation(self, runner: DummyRunner): class TestHandleInvokeAction: def test_handle_invoke_action_tool_not_present(self, runner: DummyRunner): action = AgentScratchpadUnit.Action(action_name="missing", action_input={}) - response, meta = runner._handle_invoke_action(runner.session, action, {}, []) + response, meta = runner._handle_invoke_action(action, {}, []) assert "there is not a tool named" in response - def test_tool_with_json_string_args(self, runner: DummyRunner, mocker: MockerFixture): + def test_tool_with_json_string_args(self, runner: DummyRunner): action = AgentScratchpadUnit.Action(action_name="tool", action_input=json.dumps({"a": 1})) tool_instance = MagicMock() tool_instances = {"tool": tool_instance} - mocker.patch( - "core.agent.cot_agent_runner.ToolEngine.agent_invoke", - return_value=("result", [], MagicMock(to_dict=lambda: {})), - ) + runner._tool_invoker.return_value = ("result", [], MagicMock(to_dict=lambda: {})) - response, meta = runner._handle_invoke_action(runner.session, action, tool_instances, []) + response, meta = runner._handle_invoke_action(action, tool_instances, []) assert response == "result" class TestOrganizeHistoricPromptMessages: def test_empty_history(self, runner: DummyRunner, mocker: MockerFixture): mocker.patch( - "core.agent.cot_agent_runner.AgentHistoryPromptTransform.get_prompt", + "services.agent.chat.cot_runner.AgentHistoryPromptTransform.get_prompt", return_value=[], ) result = runner._organize_historic_prompt_messages([]) @@ -241,11 +230,11 @@ def test_run_handles_empty_parser_output(self, runner: DummyRunner, mocker: Mock message.id = "msg-id" mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[], ) - results = list(runner.run(runner.session, message, "query", {})) + results = list(runner.run(message, "query", {})) assert isinstance(results, list) assert "session" not in runner.create_agent_thought.call_args.kwargs assert all("session" not in call.kwargs for call in runner.save_agent_thought.call_args_list) @@ -257,19 +246,16 @@ def test_run_with_action_and_tool_invocation(self, runner: DummyRunner, mocker: action = AgentScratchpadUnit.Action(action_name="tool", action_input={}) mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[action], ) - mocker.patch( - "core.agent.cot_agent_runner.ToolEngine.agent_invoke", - return_value=("ok", [], MagicMock(to_dict=lambda: {})), - ) + runner._tool_invoker.return_value = ("ok", [], MagicMock(to_dict=lambda: {})) runner.agent_callback = None with pytest.raises(AgentMaxIterationError): - list(runner.run(runner.session, message, "query", {"tool": MagicMock()})) + list(runner.run(message, "query", {"tool": MagicMock()})) def test_run_respects_max_iteration_boundary(self, runner: DummyRunner, mocker: MockerFixture): runner.app_config.agent.max_iteration = 1 @@ -279,30 +265,27 @@ def test_run_respects_max_iteration_boundary(self, runner: DummyRunner, mocker: action = AgentScratchpadUnit.Action(action_name="tool", action_input={}) mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[action], ) - mocker.patch( - "core.agent.cot_agent_runner.ToolEngine.agent_invoke", - return_value=("ok", [], MagicMock(to_dict=lambda: {})), - ) + runner._tool_invoker.return_value = ("ok", [], MagicMock(to_dict=lambda: {})) runner.agent_callback = None with pytest.raises(AgentMaxIterationError): - list(runner.run(runner.session, message, "query", {"tool": MagicMock()})) + list(runner.run(message, "query", {"tool": MagicMock()})) def test_run_basic_flow(self, runner: DummyRunner, mocker: MockerFixture): message = _make_message() message.id = "msg-id" mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[], ) - results = list(runner.run(runner.session, message, "query", {"name": "John"})) + results = list(runner.run(message, "query", {"name": "John"})) assert results def test_run_max_iteration_error(self, runner: DummyRunner, mocker: MockerFixture): @@ -313,12 +296,12 @@ def test_run_max_iteration_error(self, runner: DummyRunner, mocker: MockerFixtur action = AgentScratchpadUnit.Action(action_name="tool", action_input={}) mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[action], ) with pytest.raises(AgentMaxIterationError): - list(runner.run(runner.session, message, "query", {})) + list(runner.run(message, "query", {})) def test_run_increase_usage_aggregation(self, runner: DummyRunner, mocker: MockerFixture): message = _make_message() @@ -344,7 +327,7 @@ def test_run_increase_usage_aggregation(self, runner: DummyRunner, mocker: Mocke action = AgentScratchpadUnit.Action(action_name="tool", action_input={}) handle_output = mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", side_effect=[ [action], [], @@ -358,16 +341,13 @@ def _handle_side_effect(chunks, usage_dict): handle_output.side_effect = _handle_side_effect runner.model_instance.invoke_llm = MagicMock(return_value=[]) - mocker.patch( - "core.agent.cot_agent_runner.ToolEngine.agent_invoke", - return_value=("ok", [], MagicMock(to_dict=lambda: {})), - ) + runner._tool_invoker.return_value = ("ok", [], MagicMock(to_dict=lambda: {})) fake_prompt_tool = MagicMock() fake_prompt_tool.name = "tool" runner._init_prompt_tools = MagicMock(return_value=({"tool": MagicMock()}, [fake_prompt_tool])) - results = list(runner.run(runner.session, message, "query", {})) + results = list(runner.run(message, "query", {})) final_usage = results[-1].delta.usage assert final_usage is not None assert final_usage.prompt_tokens == 2 @@ -401,12 +381,12 @@ def provider_chunks(): runner.model_instance.invoke_llm.return_value = provider_chunks() mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", side_effect=lambda chunks, _usage: list(chunks), ) - results = list(runner.run(session, message, "query", {})) - assert events == ["commit", "close", "first-chunk"] + results = list(runner.run(message, "query", {})) + assert events == ["first-chunk"] assert runner.model_instance.invoke_llm.call_args.kwargs["request_metadata"] == { "app_id": "app", "app_type": CreditUsageAppType.AGENT, @@ -419,13 +399,13 @@ def test_run_usage_missing_key_branch(self, runner: DummyRunner, mocker: MockerF message.id = "msg-id" mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[], ) runner.model_instance.invoke_llm = MagicMock(return_value=[]) - list(runner.run(runner.session, message, "query", {})) + list(runner.run(message, "query", {})) def test_run_prompt_tool_update_branch(self, runner: DummyRunner, mocker: MockerFixture): message = _make_message() @@ -436,14 +416,11 @@ def test_run_prompt_tool_update_branch(self, runner: DummyRunner, mocker: Mocker # First iteration → action # Second iteration → no action (empty list) mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", side_effect=[[action], []], ) - mocker.patch( - "core.agent.cot_agent_runner.ToolEngine.agent_invoke", - return_value=("ok", [], MagicMock(to_dict=lambda: {})), - ) + runner._tool_invoker.return_value = ("ok", [], MagicMock(to_dict=lambda: {})) runner.app_config.agent.max_iteration = 5 @@ -455,7 +432,7 @@ def test_run_prompt_tool_update_branch(self, runner: DummyRunner, mocker: Mocker runner.update_prompt_message_tool = MagicMock() runner.agent_callback = None - list(runner.run(runner.session, message, "query", {})) + list(runner.run(message, "query", {})) runner.update_prompt_message_tool.assert_called_once() @@ -496,18 +473,15 @@ def test_init_react_state_resets_state(self, runner: DummyRunner, mocker: Mocker class TestHandleInvokeActionExtended: - def test_tool_with_invalid_json_string_args(self, runner: DummyRunner, mocker: MockerFixture): + def test_tool_with_invalid_json_string_args(self, runner: DummyRunner): action = AgentScratchpadUnit.Action(action_name="tool", action_input="not-json") tool_instance = MagicMock() tool_instances = {"tool": tool_instance} - mocker.patch( - "core.agent.cot_agent_runner.ToolEngine.agent_invoke", - return_value=("ok", ["file1"], MagicMock(to_dict=lambda: {"k": "v"})), - ) + runner._tool_invoker.return_value = ("ok", ["file1"], MagicMock(to_dict=lambda: {"k": "v"})) message_file_ids = [] - response, meta = runner._handle_invoke_action(runner.session, action, tool_instances, message_file_ids) + response, meta = runner._handle_invoke_action(action, tool_instances, message_file_ids) assert response == "ok" assert message_file_ids == ["file1"] @@ -538,7 +512,7 @@ def test_user_message_flushes_scratchpad(self, runner: DummyRunner, mocker: Mock runner.history_prompt_messages = [user_message] mock_transform = mocker.patch( - "core.agent.cot_agent_runner.AgentHistoryPromptTransform", + "services.agent.chat.cot_runner.AgentHistoryPromptTransform", ) mock_transform.return_value.get_prompt.return_value = ["final"] @@ -558,7 +532,7 @@ def test_agent_history_transform_invocation(self, runner: DummyRunner, mocker: M mock_transform.get_prompt.return_value = [] mocker.patch( - "core.agent.cot_agent_runner.AgentHistoryPromptTransform", + "services.agent.chat.cot_runner.AgentHistoryPromptTransform", return_value=mock_transform, ) @@ -573,11 +547,11 @@ def test_run_with_no_action_final_answer_empty(self, runner: DummyRunner, mocker message.id = "msg-id" mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=["thinking"], ) - results = list(runner.run(runner.session, message, "query", {})) + results = list(runner.run(message, "query", {})) assert any(hasattr(r, "delta") for r in results) def test_run_with_final_answer_action_string(self, runner: DummyRunner, mocker: MockerFixture): @@ -587,11 +561,11 @@ def test_run_with_final_answer_action_string(self, runner: DummyRunner, mocker: action = AgentScratchpadUnit.Action(action_name="Final Answer", action_input="done") mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[action], ) - results = list(runner.run(runner.session, message, "query", {})) + results = list(runner.run(message, "query", {})) assert results[-1].delta.message.content == "done" def test_run_with_final_answer_action_dict(self, runner: DummyRunner, mocker: MockerFixture): @@ -601,11 +575,11 @@ def test_run_with_final_answer_action_dict(self, runner: DummyRunner, mocker: Mo action = AgentScratchpadUnit.Action(action_name="Final Answer", action_input={"a": 1}) mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[action], ) - results = list(runner.run(runner.session, message, "query", {})) + results = list(runner.run(message, "query", {})) assert json.loads(results[-1].delta.message.content) == {"a": 1} def test_run_with_string_final_answer(self, runner: DummyRunner, mocker: MockerFixture): @@ -616,9 +590,9 @@ def test_run_with_string_final_answer(self, runner: DummyRunner, mocker: MockerF action = AgentScratchpadUnit.Action(action_name="Final Answer", action_input="12345") mocker.patch( - "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", + "services.agent.chat.cot_runner.CotAgentOutputParser.handle_react_stream_output", return_value=[action], ) - results = list(runner.run(runner.session, message, "query", {})) + results = list(runner.run(message, "query", {})) assert results[-1].delta.message.content == "12345" diff --git a/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py index 8e7093fd123a1d..1213e5464cff5d 100644 --- a/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py @@ -3,8 +3,8 @@ import pytest from pytest_mock import MockerFixture -from core.agent.cot_chat_agent_runner import CotChatAgentRunner from graphon.model_runtime.entities.message_entities import TextPromptMessageContent +from services.agent.chat.cot_chat_runner import CotChatAgentRunner from tests.unit_tests.core.agent.conftest import ( DummyAgentConfig, DummyAppConfig, @@ -61,7 +61,7 @@ def test_organize_system_prompt_success(self, runner, mocker: MockerFixture): runner.app_config = DummyAppConfig(DummyAgentConfig(DummyPrompt(first_prompt))) mocker.patch( - "core.agent.cot_chat_agent_runner.jsonable_encoder", + "services.agent.chat.cot_chat_runner.jsonable_encoder", return_value=[{"name": "tool1"}, {"name": "tool2"}], ) @@ -91,8 +91,8 @@ def test_organize_user_query_no_files(self, runner, files): assert len(result) == 1 assert result[0].content == "query" - @patch("core.agent.cot_chat_agent_runner.UserPromptMessage") - @patch("core.agent.cot_chat_agent_runner.file_manager.to_prompt_message_content") + @patch("services.agent.chat.cot_chat_runner.UserPromptMessage") + @patch("services.agent.chat.cot_chat_runner.file_manager.to_prompt_message_content") def test_organize_user_query_with_image_file_default_config(self, mock_to_prompt, mock_user_prompt, runner): from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent @@ -116,8 +116,8 @@ def test_organize_user_query_with_image_file_default_config(self, mock_to_prompt image_detail_config=ImagePromptMessageContent.DETAIL.LOW, ) - @patch("core.agent.cot_chat_agent_runner.UserPromptMessage") - @patch("core.agent.cot_chat_agent_runner.file_manager.to_prompt_message_content") + @patch("services.agent.chat.cot_chat_runner.UserPromptMessage") + @patch("services.agent.chat.cot_chat_runner.file_manager.to_prompt_message_content") def test_organize_user_query_with_image_file_high_detail(self, mock_to_prompt, mock_user_prompt, runner): from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent @@ -143,7 +143,7 @@ def test_organize_user_query_with_image_file_high_detail(self, mock_to_prompt, m image_detail_config=ImagePromptMessageContent.DETAIL.HIGH, ) - @patch("core.agent.cot_chat_agent_runner.file_manager.to_prompt_message_content") + @patch("services.agent.chat.cot_chat_runner.file_manager.to_prompt_message_content") def test_organize_user_query_with_text_file_no_config(self, mock_to_prompt, runner): mock_to_prompt.return_value = TextPromptMessageContent(data="file_content") runner.files = ["file1"] diff --git a/api/tests/unit_tests/core/agent/test_cot_completion_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_completion_agent_runner.py index e910c59d3687bb..4a67dc39521f27 100644 --- a/api/tests/unit_tests/core/agent/test_cot_completion_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_cot_completion_agent_runner.py @@ -3,13 +3,13 @@ import pytest from pytest_mock import MockerFixture -from core.agent.cot_completion_agent_runner import CotCompletionAgentRunner from graphon.model_runtime.entities.message_entities import ( AssistantPromptMessage, ImagePromptMessageContent, TextPromptMessageContent, UserPromptMessage, ) +from services.agent.chat.cot_completion_runner import CotCompletionAgentRunner # ----------------------------- # Fixtures @@ -26,7 +26,7 @@ def runner(mocker: MockerFixture, dummy_tool_factory): runner._agent_scratchpad = [] mocker.patch( - "core.agent.cot_completion_agent_runner.jsonable_encoder", + "services.agent.chat.cot_completion_runner.jsonable_encoder", side_effect=lambda tools: [{"name": t.name} for t in tools], ) diff --git a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py index 0b6e9bd429ede0..097a23abcfac43 100644 --- a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py @@ -4,7 +4,7 @@ from datetime import UTC, datetime from decimal import Decimal from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, create_autospec import pytest from pytest_mock import MockerFixture @@ -14,8 +14,6 @@ from core.agent.entities import AgentEntity from core.agent.errors import AgentMaxIterationError -from core.agent.fc_agent_runner import FunctionCallAgentRunner -from core.app.app_config.entities import DatasetRetrieveConfigEntity from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig from core.app.apps.base_app_queue_manager import PublishFrom from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager @@ -23,6 +21,7 @@ from core.app.entities.queue_entities import QueueMessageFileEvent from core.credit_usage import CreditUsageAppType from core.model_manager import ModelInstance +from extensions.application_services.workflow import WorkflowExecutionDependencies from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.message_entities import ( AssistantPromptMessage, @@ -36,21 +35,22 @@ from libs.datetime_utils import naive_utc_now from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus from models.model import AppMode, Conversation, Message, StorageType, UploadFile -from services.workflow_run_agg import WorkflowRunAgg +from repositories.app.generation_repository import AppGenerationRepository +from services.agent.chat.function_call_runner import FunctionCallAgentRunner # ============================== # Dummy Helper Classes # ============================== -def build_usage(pt=1, ct=1, tt=2) -> LLMUsage: +def build_usage(pt: int = 1, ct: int = 1, tt: int = 2) -> LLMUsage: usage = LLMUsage.empty_usage() usage.prompt_tokens = pt usage.completion_tokens = ct usage.total_tokens = tt - usage.prompt_price = 0 - usage.completion_price = 0 - usage.total_price = 0 + usage.prompt_price = Decimal(0) + usage.completion_price = Decimal(0) + usage.total_price = Decimal(0) return usage @@ -85,8 +85,8 @@ def _make_message(*, message_id: str = "m1", conversation_id: str = "conv1") -> class DummyMessage: - def __init__(self, content: str | None = None, tool_calls: list[Any] | None = None): - self.content: str | None = content + def __init__(self, content: object = None, tool_calls: list[Any] | None = None) -> None: + self.content: object = content self.tool_calls: list[Any] = tool_calls or [] @@ -97,7 +97,7 @@ class DummyDelta: class DummyChunk: - def __init__(self, message: DummyMessage | None = None, usage: LLMUsage | None = None): + def __init__(self, message: DummyMessage | None = None, usage: LLMUsage | None = None) -> None: self.delta: DummyDelta = DummyDelta(message=message, usage=usage) @@ -107,13 +107,26 @@ def __init__( message: DummyMessage | None = None, usage: LLMUsage | None = None, prompt_messages: list[DummyMessage] | None = None, - ): + ) -> None: self.message: DummyMessage | None = message self.usage: LLMUsage | None = usage self.prompt_messages: list[DummyMessage] = prompt_messages or [] self.system_fingerprint: str = "" +class _TestFunctionCallAgentRunner(FunctionCallAgentRunner): + """Expose dynamically mocked collaborators used by focused runner tests.""" + + session: Session + model_instance: Any + create_agent_thought: Any + save_agent_thought: Any + queue_manager: Any + _init_prompt_tools: Any + _tool_invoker: Any + app_config: Any + + # ============================== # Fixtures # ============================== @@ -121,7 +134,9 @@ def __init__( @pytest.fixture def real_runner( - sqlite_session: Session, + app_records: AppGenerationRepository, + workflow_runtime: WorkflowExecutionDependencies, + agent_tool_invoker: MagicMock, agent_records: tuple[Conversation, Message], agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, @@ -133,8 +148,10 @@ def real_runner( assert app_config.agent is not None app_config.agent.strategy = AgentEntity.Strategy.FUNCTION_CALLING return FunctionCallAgentRunner( - execution_driver=WorkflowRunAgg.run, - session=sqlite_session, + records=app_records, + dataset_tools=workflow_runtime.dataset_tools, + tool_invoker=agent_tool_invoker, + workflow_runtime=workflow_runtime, tenant_id=app_config.tenant_id, application_generate_entity=agent_generate_entity, conversation=conversation, @@ -160,16 +177,22 @@ def _tool_call(call_id: str, name: str, arguments: str) -> AssistantPromptMessag @pytest.fixture -def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCallAgentRunner]: +def runner( + mocker: MockerFixture, + sqlite_engine: Engine, + app_records: AppGenerationRepository, + workflow_runtime: WorkflowExecutionDependencies, + agent_tool_invoker: MagicMock, +) -> Iterator[_TestFunctionCallAgentRunner]: # Completely bypass BaseAgentRunner __init__ to avoid DB / Flask context mocker.patch( - "core.agent.base_agent_runner.BaseAgentRunner.__init__", + "services.agent.chat.base_runner.BaseAgentRunner.__init__", return_value=None, ) # Patch streaming chunk models to avoid validation on dummy message objects - mocker.patch("core.agent.fc_agent_runner.LLMResultChunk", MagicMock) - mocker.patch("core.agent.fc_agent_runner.LLMResultChunkDelta", MagicMock) + mocker.patch("services.agent.chat.function_call_runner.LLMResultChunk", MagicMock) + mocker.patch("services.agent.chat.function_call_runner.LLMResultChunkDelta", MagicMock) app_config = MagicMock() app_config.app_id = "app" @@ -191,8 +214,11 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCal message = _make_message(message_id="msg1") conversation = _make_conversation() - runner = FunctionCallAgentRunner( - execution_driver=WorkflowRunAgg.run, + runner = _TestFunctionCallAgentRunner( + records=app_records, + dataset_tools=workflow_runtime.dataset_tools, + tool_invoker=agent_tool_invoker, + workflow_runtime=workflow_runtime, tenant_id="tenant", application_generate_entity=application_generate_entity, conversation=conversation, @@ -216,6 +242,7 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCal runner.message = message runner.user_id = "user" runner.model_instance = model_instance + runner._tool_invoker = agent_tool_invoker runner.stream_tool_call = False runner.memory = None @@ -233,6 +260,7 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCal runner.update_prompt_message_tool = MagicMock() try: + runner._records = app_records yield runner finally: runner.session.close() @@ -311,7 +339,7 @@ def test_extract_blocking_tool_calls(self, real_runner: FunctionCallAgentRunner) class TestInitSystemMessage: - def test_init_system_message_empty_prompt_messages(self, runner: FunctionCallAgentRunner): + def test_init_system_message_empty_prompt_messages(self, runner: _TestFunctionCallAgentRunner) -> None: result = runner._init_system_message("system", []) assert len(result) == 1 @@ -320,7 +348,7 @@ def test_init_system_message_insert_at_start(self, real_runner: FunctionCallAgen result = real_runner._init_system_message("system", [query]) assert result == [SystemPromptMessage(content="system"), query] - def test_init_system_message_no_template(self, runner: FunctionCallAgentRunner): + def test_init_system_message_no_template(self, runner: _TestFunctionCallAgentRunner) -> None: result = runner._init_system_message("", []) assert result == [] @@ -331,18 +359,20 @@ def test_init_system_message_no_template(self, runner: FunctionCallAgentRunner): class TestOrganizeUserQuery: - def test_without_files(self, runner: FunctionCallAgentRunner): + def test_without_files(self, runner: _TestFunctionCallAgentRunner) -> None: result = runner._organize_user_query("query", []) assert len(result) == 1 - def test_with_none_query(self, runner: FunctionCallAgentRunner): - result = runner._organize_user_query(None, []) + def test_with_none_query(self, runner: _TestFunctionCallAgentRunner) -> None: + result = runner._organize_user_query(None, []) # pyrefly: ignore[bad-argument-type] assert len(result) == 1 - def test_with_files_uses_image_detail_config(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): + def test_with_files_uses_image_detail_config( + self, runner: _TestFunctionCallAgentRunner, mocker: MockerFixture + ) -> None: file_content = TextPromptMessageContent(data="file-content") mock_to_prompt = mocker.patch( - "core.agent.fc_agent_runner.file_manager.to_prompt_message_content", + "services.agent.chat.function_call_runner.file_manager.to_prompt_message_content", return_value=file_content, ) @@ -372,7 +402,7 @@ def test_clear_text_and_image_content(self, real_runner: FunctionCallAgentRunner assert result == [UserPromptMessage(content="hello\n[image]")] assert user_msg.content == [text, image] - def test_clear_includes_file_placeholder(self, runner: FunctionCallAgentRunner): + def test_clear_includes_file_placeholder(self, runner: _TestFunctionCallAgentRunner) -> None: text = TextPromptMessageContent(data="hello") image = ImagePromptMessageContent(format="url", mime_type="image/png") document = DocumentPromptMessageContent(format="url", mime_type="application/pdf") @@ -383,7 +413,7 @@ def test_clear_includes_file_placeholder(self, runner: FunctionCallAgentRunner): assert result[0].content == "hello\n[image]\n[file]" - def test_keeps_knowledge_retrieval_image_message(self, runner: FunctionCallAgentRunner): + def test_keeps_knowledge_retrieval_image_message(self, runner: _TestFunctionCallAgentRunner) -> None: text = TextPromptMessageContent(data="query") image = ImagePromptMessageContent(format="url", mime_type="image/png") user_msg = UserPromptMessage(name="knowledge_retrieval", content=[image, text]) @@ -399,33 +429,28 @@ def test_keeps_knowledge_retrieval_image_message(self, runner: FunctionCallAgent class TestBuildDatasetToolImageContents: - def test_returns_empty_when_vision_disabled(self, runner: FunctionCallAgentRunner): + def test_returns_empty_when_vision_disabled(self, runner: _TestFunctionCallAgentRunner) -> None: tool = MagicMock() tool.__class__.__name__ = "DatasetRetrieverTool" response = "![image](http://localhost:5001/files/890985e9-c2f1-484e-bc7b-62010a337e6d/file-preview)" - assert runner._build_dataset_tool_image_contents(runner.session, response, tool) == [] + assert runner._build_dataset_tool_image_contents(response, tool) == [] def test_builds_image_contents_from_dataset_tool_preview_links( - self, runner: FunctionCallAgentRunner, mocker: MockerFixture - ): - from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool + self, runner: _TestFunctionCallAgentRunner, mocker: MockerFixture + ) -> None: + from services.tools.dataset.tool import DatasetRetrieverTool runner.vision_enabled = True image_content = ImagePromptMessageContent(format="url", mime_type="image/png") to_prompt_content = mocker.patch( - "core.agent.fc_agent_runner.file_manager.to_prompt_message_content", + "services.agent.chat.function_call_runner.file_manager.to_prompt_message_content", return_value=image_content, ) - grant_access = mocker.patch("core.agent.fc_agent_runner.grant_upload_file_access") - sign_preview = mocker.patch( - "core.agent.fc_agent_runner.sign_upload_file_preview_url", - return_value="http://localhost:5001/files/file-id/file-preview?sign=1", - ) - build_reference = mocker.patch("core.agent.fc_agent_runner.build_file_reference", return_value="file-ref") + grant_access = mocker.patch("services.agent.chat.function_call_runner.grant_upload_file_access") upload_file = UploadFile( - tenant_id="00000000-0000-0000-0000-000000000001", + tenant_id=runner.tenant_id, storage_type=StorageType.LOCAL, key="image_files/chart.png", name="chart.png", @@ -462,10 +487,12 @@ def test_builds_image_contents_from_dataset_tool_preview_links( "file ![file](http://localhost:5001/files/11111111-1111-1111-1111-111111111111/file-preview)" ) + from core.app.app_config.entities import DatasetRetrieveConfigEntity + from core.app.entities.app_invoke_entities import InvokeFrom from core.tools.__base.tool_runtime import ToolRuntime from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ToolDescription, ToolEntity, ToolIdentity - from core.tools.utils.dataset_retriever.dataset_retriever_tool import DatasetRetrieverTool as RetrievalTool + from services.knowledge.retrieval.dataset_retrieval import DatasetRetrieval tool = DatasetRetrieverTool( entity=ToolEntity( @@ -476,25 +503,22 @@ def test_builds_image_contents_from_dataset_tool_preview_links( description=ToolDescription(human=I18nObject(en_US="Retrieve dataset"), llm="Retrieve dataset"), ), runtime=ToolRuntime(tenant_id=upload_file.tenant_id), - retrieval_tool=RetrievalTool( - tenant_id=upload_file.tenant_id, - dataset_id="dataset-id", - return_resource=True, - retriever_from="dev", - retrieve_config=DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE - ), - inputs={}, - ), + retrieval=create_autospec(DatasetRetrieval, instance=True), + dataset_id="dataset-id", + config=DatasetRetrieveConfigEntity(retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE), + top_k=2, + inputs={}, + invoke_from=InvokeFrom.DEBUGGER, + return_resource=True, + hit_callback=MagicMock(), + app_id="app", ) - contents = runner._build_dataset_tool_image_contents(session, response, tool) + contents = runner._build_dataset_tool_image_contents(response, tool) assert contents == [image_content] assert contents[0].type == PromptMessageContentType.IMAGE grant_access.assert_called_once() assert list(grant_access.call_args.args[0]) == ["890985e9-c2f1-484e-bc7b-62010a337e6d"] - sign_preview.assert_called_once_with(upload_file.id, upload_file.extension) - build_reference.assert_called_once_with(record_id=str(upload_file.id)) to_prompt_content.assert_called_once() @@ -504,14 +528,14 @@ def test_builds_image_contents_from_dataset_tool_preview_links( class TestRunMethod: - def test_run_non_streaming_no_tool_calls(self, runner: FunctionCallAgentRunner): + def test_run_non_streaming_no_tool_calls(self, runner: _TestFunctionCallAgentRunner) -> None: message = _make_message() dummy_message = DummyMessage(content="hello") result = DummyResult(message=dummy_message, usage=build_usage()) runner.model_instance.invoke_llm.return_value = result - outputs = list(runner.run(runner.session, message, "query")) + outputs = list(runner.run(message, "query")) assert len(outputs) == 1 assert "session" not in runner.create_agent_thought.call_args.kwargs assert "session" not in runner.save_agent_thought.call_args.kwargs @@ -525,7 +549,7 @@ def test_run_non_streaming_no_tool_calls(self, runner: FunctionCallAgentRunner): queue_calls = runner.queue_manager.publish.call_args_list assert any(call.args and call.args[0].__class__.__name__ == "QueueMessageEndEvent" for call in queue_calls) - def test_run_streaming_branch(self, runner: FunctionCallAgentRunner): + def test_run_streaming_branch(self, runner: _TestFunctionCallAgentRunner) -> None: runner.stream_tool_call = True events: list[str] = [] session = runner.session @@ -547,17 +571,17 @@ def record_detach(_session: Session, instance: object) -> None: content = [TextPromptMessageContent(data="hi")] chunk = DummyChunk(message=DummyMessage(content=content), usage=build_usage()) - def generator(): + def generator() -> Iterator[DummyChunk]: events.append("first-chunk") yield chunk runner.model_instance.invoke_llm.return_value = generator() - outputs = list(runner.run(session, message, "query")) - assert events == ["commit", "close", "first-chunk"] + outputs = list(runner.run(message, "query")) + assert events == ["first-chunk"] assert len(outputs) == 1 - def test_run_streaming_tool_calls_list_content(self, runner: FunctionCallAgentRunner): + def test_run_streaming_tool_calls_list_content(self, runner: _TestFunctionCallAgentRunner) -> None: message = _make_message() runner.stream_tool_call = True @@ -569,7 +593,7 @@ def test_run_streaming_tool_calls_list_content(self, runner: FunctionCallAgentRu content = [TextPromptMessageContent(data="hi")] chunk = DummyChunk(message=DummyMessage(content=content, tool_calls=[tool_call]), usage=build_usage()) - def generator(): + def generator() -> Iterator[DummyChunk]: yield chunk final_message = DummyMessage(content="done", tool_calls=[]) @@ -577,10 +601,10 @@ def generator(): runner.model_instance.invoke_llm.side_effect = [generator(), final_result] - outputs = list(runner.run(runner.session, message, "query")) + outputs = list(runner.run(message, "query")) assert len(outputs) >= 1 - def test_run_non_streaming_list_content(self, runner: FunctionCallAgentRunner): + def test_run_non_streaming_list_content(self, runner: _TestFunctionCallAgentRunner) -> None: message = _make_message() content = [TextPromptMessageContent(data="hi")] dummy_message = DummyMessage(content=content) @@ -588,11 +612,13 @@ def test_run_non_streaming_list_content(self, runner: FunctionCallAgentRunner): runner.model_instance.invoke_llm.return_value = result - outputs = list(runner.run(runner.session, message, "query")) + outputs = list(runner.run(message, "query")) assert len(outputs) == 1 assert runner.save_agent_thought.call_args.kwargs["thought"] == "hi" - def test_run_streaming_tool_call_inputs_type_error(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): + def test_run_streaming_tool_call_inputs_type_error( + self, runner: _TestFunctionCallAgentRunner, mocker: MockerFixture + ) -> None: message = _make_message() runner.stream_tool_call = True @@ -603,24 +629,24 @@ def test_run_streaming_tool_call_inputs_type_error(self, runner: FunctionCallAge chunk = DummyChunk(message=DummyMessage(content="hi", tool_calls=[tool_call]), usage=build_usage()) - def generator(): + def generator() -> Iterator[DummyChunk]: yield chunk runner.model_instance.invoke_llm.return_value = generator() real_dumps = json.dumps - def flaky_dumps(obj, *args, **kwargs): - if kwargs.get("ensure_ascii") is False: - return real_dumps(obj, *args, **kwargs) + def flaky_dumps(obj: object, *, ensure_ascii: bool = True) -> str: + if not ensure_ascii: + return real_dumps(obj, ensure_ascii=ensure_ascii) raise TypeError("boom") - mocker.patch("core.agent.fc_agent_runner.json.dumps", side_effect=flaky_dumps) + mocker.patch("services.agent.chat.function_call_runner.json.dumps", side_effect=flaky_dumps) - outputs = list(runner.run(runner.session, message, "query")) + outputs = list(runner.run(message, "query")) assert len(outputs) == 1 - def test_run_with_missing_tool_instance(self, runner: FunctionCallAgentRunner): + def test_run_with_missing_tool_instance(self, runner: _TestFunctionCallAgentRunner) -> None: message = _make_message() tool_call = MagicMock() @@ -635,10 +661,10 @@ def test_run_with_missing_tool_instance(self, runner: FunctionCallAgentRunner): runner.model_instance.invoke_llm.side_effect = [result, final_result] - outputs = list(runner.run(runner.session, message, "query")) + outputs = list(runner.run(message, "query")) assert len(outputs) >= 1 - def test_run_with_tool_instance_and_files(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): + def test_run_with_tool_instance_and_files(self, runner: _TestFunctionCallAgentRunner) -> None: message = _make_message() tool_call = MagicMock() @@ -659,12 +685,9 @@ def test_run_with_tool_instance_and_files(self, runner: FunctionCallAgentRunner, tool_invoke_meta = MagicMock() tool_invoke_meta.to_dict.return_value = {"ok": True} - mocker.patch( - "core.agent.fc_agent_runner.ToolEngine.agent_invoke", - return_value=("ok", ["file1"], tool_invoke_meta), - ) + runner._tool_invoker.return_value = ("ok", ["file1"], tool_invoke_meta) - outputs = list(runner.run(runner.session, message, "query")) + outputs = list(runner.run(message, "query")) assert len(outputs) >= 1 assert any( isinstance(call.args[0], QueueMessageFileEvent) @@ -673,7 +696,7 @@ def test_run_with_tool_instance_and_files(self, runner: FunctionCallAgentRunner, for call in runner.queue_manager.publish.call_args_list ) - def test_run_max_iteration_error(self, runner: FunctionCallAgentRunner): + def test_run_max_iteration_error(self, runner: _TestFunctionCallAgentRunner) -> None: runner.app_config.agent.max_iteration = 0 message = _make_message() @@ -689,4 +712,4 @@ def test_run_max_iteration_error(self, runner: FunctionCallAgentRunner): runner.model_instance.invoke_llm.return_value = result with pytest.raises(AgentMaxIterationError): - list(runner.run(runner.session, message, "query")) + list(runner.run(message, "query")) 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 5346080c6e6fe5..81cf54b7d4bc96 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 @@ -16,26 +16,16 @@ from pytest_mock import MockerFixture from sqlalchemy import Engine as _app_generator_Engine from sqlalchemy import event as sqlalchemy_event -from sqlalchemy.engine import Engine as _advanced_chat_app_generator_Engine +from sqlalchemy import inspect from sqlalchemy.orm import Session -import core.app.apps.advanced_chat.app_generator as app_generator_module import models.human_input_entities # noqa: F401 from constants import UUID_NIL from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat import app_generator as adv_app_gen_module -from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator -from core.app.apps.advanced_chat.generate_task_pipeline import ( - ConversationSnapshot, - MessageSnapshot, - WorkflowSnapshot, -) from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.workflow import app_generator as wf_app_gen_module from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom from core.app.file_access import DatabaseFileAccessController -from core.app.task_pipeline import message_cycle_manager -from core.app.task_pipeline.message_cycle_manager import MessageCycleManager from core.app.workflow.file_runtime import DifyWorkflowFileRuntime from core.ops.ops_trace_manager import TraceQueueManager from core.workflow import node_factory as node_factory_module @@ -72,7 +62,14 @@ from models.workflow import Workflow as _app_generator_Workflow from models.workflow import Workflow as _pause_resume_Workflow from models.workflow import WorkflowType +from services.app.generation.adapters import message_cycle as message_cycle_manager +from services.app.generation.adapters.message_cycle import MessageCycleManager +from services.app.generation.ports import ConversationSnapshot, MessageSnapshot, WorkflowSnapshot from services.errors.conversation import ConversationNotExistsError +from services.workflow.execution.adapters.chatflow import app_generator as adv_app_gen_module +from services.workflow.execution.adapters.chatflow.app_generator import AdvancedChatAppGenerator +from services.workflow.execution.adapters.chatflow.record_preparation import prepare_chat_records +from services.workflow.execution.ports import WorkflowRuntime from services.workflow_run_agg import WorkflowRunAgg from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.model_factories import make_upload_file @@ -158,8 +155,10 @@ def _make_message( class TestAdvancedChatAppGeneratorValidation: - def test_generate_requires_query(self, unbound_session: Session): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_generate_requires_query(self, unbound_session: Session, *, workflow_runtime: WorkflowRuntime): + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) with pytest.raises(ValueError, match="query is required"): generator.generate( @@ -170,11 +169,12 @@ def test_generate_requires_query(self, unbound_session: Session): invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", streaming=False, - session=unbound_session, ) - def test_generate_requires_string_query(self, unbound_session: Session): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_generate_requires_string_query(self, unbound_session: Session, *, workflow_runtime: WorkflowRuntime): + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) with pytest.raises(ValueError, match="query must be a string"): generator.generate( @@ -185,11 +185,14 @@ def test_generate_requires_string_query(self, unbound_session: Session): invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", streaming=False, - session=unbound_session, ) - def test_single_iteration_generate_validates_args(self, unbound_session: Session): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_single_iteration_generate_validates_args( + self, unbound_session: Session, *, workflow_runtime: WorkflowRuntime + ): + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) with pytest.raises(ValueError, match="node_id is required"): generator.single_iteration_generate( @@ -199,7 +202,6 @@ def test_single_iteration_generate_validates_args(self, unbound_session: Session user=_app_generator_make_account(), args={"inputs": {}}, streaming=False, - session=unbound_session, ) with pytest.raises(ValueError, match="inputs is required"): @@ -210,11 +212,12 @@ def test_single_iteration_generate_validates_args(self, unbound_session: Session user=_app_generator_make_account(), args={}, streaming=False, - session=unbound_session, ) - def test_single_loop_generate_validates_args(self, unbound_session: Session): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_single_loop_generate_validates_args(self, unbound_session: Session, *, workflow_runtime: WorkflowRuntime): + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) with pytest.raises(ValueError, match="node_id is required"): generator.single_loop_generate( @@ -222,9 +225,8 @@ def test_single_loop_generate_validates_args(self, unbound_session: Session): workflow=_app_generator_make_workflow(), node_id="", user=_app_generator_make_account(), - args=SimpleNamespace(inputs={}), + args={"inputs": {}}, streaming=False, - session=unbound_session, ) with pytest.raises(ValueError, match="inputs is required"): @@ -233,9 +235,8 @@ def test_single_loop_generate_validates_args(self, unbound_session: Session): workflow=_app_generator_make_workflow(), node_id="node", user=_app_generator_make_account(), - args=SimpleNamespace(inputs=None), + args={"inputs": None}, streaming=False, - session=unbound_session, ) @@ -252,46 +253,41 @@ def _build_app_config() -> WorkflowUIBasedAppConfig: ) def test_generate_loads_conversation_and_files( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) app_config = self._build_app_config() conversation = _make_conversation() built_files: list[object] = [] build_files_called = {"called": False} captured: dict[str, object] = {} - session = unbound_session get_conversation = MagicMock(return_value=conversation) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.ConversationService.get_conversation", + workflow_runtime.chat_records, + "conversation", get_conversation, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", + "services.workflow.execution.adapters.chatflow.app_generator.FileUploadConfigManager.convert", lambda *args, **kwargs: {"enabled": True}, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.file_factory.build_from_mappings", + "services.workflow.execution.adapters.chatflow.app_generator.file_factory.build_from_mappings", lambda **kwargs: build_files_called.update({"called": True}) or built_files, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppConfigManager.get_app_config", lambda **kwargs: app_config, ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=unbound_session), - ) monkeypatch.setattr(generator, "_prepare_user_inputs", lambda **kwargs: kwargs["user_inputs"]) dummy_trace_queue_manager = type( @@ -303,7 +299,10 @@ def test_generate_loads_conversation_and_files( ) }, ) - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.TraceQueueManager", dummy_trace_queue_manager) + monkeypatch.setattr( + "services.workflow.execution.adapters.chatflow.app_generator.TraceQueueManager", + dummy_trace_queue_manager, + ) def _fake_generate(**kwargs): captured.update(kwargs) @@ -329,23 +328,28 @@ def _fake_generate(**kwargs): invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", streaming=False, - session=session, ) assert result == {"ok": True} assert captured["conversation"] is conversation assert captured["application_generate_entity"].files == built_files - assert captured["session"] is session assert build_files_called["called"] is True - assert get_conversation.call_args.kwargs["session"] is session - def test_resume_delegates_to_generate(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_resume_delegates_to_generate( + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + *, + workflow_runtime: WorkflowRuntime, + ): + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) existing_trace_manager = SimpleNamespace(app_id="existing-app", user_id="existing-user") application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( task_id="task", app_config=self._build_app_config(), - inputs={}, + _inputs={}, query="hello", files=[], user_id="user", @@ -373,7 +377,6 @@ def _fake_generate(**kwargs): user=_make_end_user(), conversation=_make_conversation(), message=_make_message(), - session=unbound_session, application_generate_entity=application_generate_entity, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), @@ -387,43 +390,33 @@ def _fake_generate(**kwargs): assert captured_graph_runtime_state is not None def test_single_iteration_generate_builds_debug_task( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) app_config = self._build_app_config() captured: dict[str, object] = {} prefill_calls: list[object] = [] - draft_sessions: list[object] = [] var_loader = SimpleNamespace(loader="draft") workflow = _app_generator_make_workflow() - session = unbound_session monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppConfigManager.get_app_config", lambda **kwargs: app_config, ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: SimpleNamespace(repo="execution"), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: SimpleNamespace(repo="node"), - ) - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.DraftVarLoader", lambda **kwargs: var_loader) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=unbound_session), - ) - - class _DraftVarService: - def __init__(self, session): - draft_sessions.append(session) - def prefill_conversation_variable_default_values(self, workflow, user_id): - prefill_calls.append((workflow, user_id)) + def load_variables(workflow, user_id): + prefill_calls.append((workflow, user_id)) + return var_loader - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.WorkflowDraftVariableService", _DraftVarService) + generator = AdvancedChatAppGenerator( + draft_variable_loader=load_variables, + runtime=workflow_runtime, + execution_driver=workflow_runtime.execution_driver, + ) def _fake_generate(**kwargs): captured.update(kwargs) @@ -438,55 +431,42 @@ def _fake_generate(**kwargs): user=_app_generator_make_account(), args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, streaming=False, - session=session, ) assert result == {"ok": True} assert prefill_calls == [(workflow, "user-id")] - assert draft_sessions == [session] assert captured["variable_loader"] is var_loader - assert captured["session"] is session assert captured["application_generate_entity"].single_iteration_run.node_id == "node-1" assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" def test_single_loop_generate_builds_debug_task( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) app_config = self._build_app_config() captured: dict[str, object] = {} prefill_calls: list[object] = [] - draft_sessions: list[object] = [] var_loader = SimpleNamespace(loader="draft") workflow = _app_generator_make_workflow() - session = unbound_session monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppConfigManager.get_app_config", lambda **kwargs: app_config, ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: SimpleNamespace(repo="execution"), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: SimpleNamespace(repo="node"), - ) - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.DraftVarLoader", lambda **kwargs: var_loader) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=unbound_session), - ) - class _DraftVarService: - def __init__(self, session): - draft_sessions.append(session) + def load_variables(workflow, user_id): + prefill_calls.append((workflow, user_id)) + return var_loader - def prefill_conversation_variable_default_values(self, workflow, user_id): - prefill_calls.append((workflow, user_id)) - - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.WorkflowDraftVariableService", _DraftVarService) + generator = AdvancedChatAppGenerator( + draft_variable_loader=load_variables, + runtime=workflow_runtime, + execution_driver=workflow_runtime.execution_driver, + ) def _fake_generate(**kwargs): captured.update(kwargs) @@ -499,23 +479,27 @@ def _fake_generate(**kwargs): workflow=workflow, node_id="node-2", user=_app_generator_make_account(), - args=SimpleNamespace(inputs={"foo": "bar"}, trace_session_id="session-1"), + args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, streaming=False, - session=session, ) assert result == {"ok": True} assert prefill_calls == [(workflow, "user-id")] - assert draft_sessions == [session] assert captured["variable_loader"] is var_loader - assert captured["session"] is session assert captured["application_generate_entity"].single_loop_run.node_id == "node-2" assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" def test_generate_internal_flow_initial_conversation_with_pause_layer( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 0 app_config = self._build_app_config() @@ -549,21 +533,21 @@ def _record_commit(session: Session) -> None: captured: dict[str, object] = {} thread_data: dict[str, object] = {} init_records = MagicMock(return_value=(conversation, message)) - get_thread_messages_length = MagicMock(return_value=2) + get_thread_messages_length = MagicMock(return_value=3) - monkeypatch.setattr(generator, "_init_generate_records", init_records) + monkeypatch.setattr(generator._runtime.chat_records, "initialize", init_records) + monkeypatch.setattr(generator._runtime.chat_records, "dialogue_count", get_thread_messages_length) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.get_thread_messages_length", get_thread_messages_length - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.MessageBasedAppQueueManager", + "services.workflow.execution.adapters.chatflow.app_generator.MessageBasedAppQueueManager", lambda **kwargs: SimpleNamespace(**kwargs), ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.current_app", + "services.workflow.execution.adapters.chatflow.app_generator.current_app", SimpleNamespace(_get_current_object=lambda: SimpleNamespace(name="flask")), ) - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.contextvars.copy_context", lambda: "ctx") + monkeypatch.setattr( + "services.workflow.execution.adapters.chatflow.app_generator.contextvars.copy_context", lambda: "ctx" + ) class _Thread: def __init__(self, *, target, kwargs): @@ -580,11 +564,7 @@ def join(self, 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=sqlite_engine, session=sqlite_session), - ) + monkeypatch.setattr("services.workflow.execution.adapters.chatflow.app_generator.threading.Thread", _Thread) monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory") monkeypatch.setattr( generator, @@ -592,7 +572,7 @@ def is_alive(self): lambda **kwargs: captured.update(kwargs) or {"raw": True}, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppGenerateResponseConverter.convert", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppGenerateResponseConverter.convert", lambda response, invoke_from: {"response": response, "invoke_from": invoke_from}, ) @@ -603,7 +583,6 @@ def is_alive(self): user=_app_generator_make_account(account_id="user"), invoke_from=InvokeFrom.WEB_APP, application_generate_entity=application_generate_entity, - session=sqlite_session, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), conversation=None, @@ -618,19 +597,25 @@ def is_alive(self): assert thread_data["join_timeout"] == 300 assert thread_data["kwargs"]["pause_state_config"] is pause_state_config assert generator._dialogue_count == 3 - assert init_records.call_args.kwargs["session"] is sqlite_session - get_thread_messages_length.assert_called_once_with(conversation.id, session=sqlite_session) - assert commit_count == 1 - assert json.loads(conversation.override_model_configs) == {"feature": True} + get_thread_messages_length.assert_called_once_with(conversation.id) + assert commit_count == 0 + assert init_records.call_args.kwargs["seed"].conversation["override_model_configs"] == workflow.features assert captured["draft_var_saver_factory"] == "draft-factory" assert isinstance(captured["workflow"], WorkflowSnapshot) assert isinstance(captured["conversation"], ConversationSnapshot) assert isinstance(captured["message"], MessageSnapshot) def test_generate_internal_flow_with_existing_records_skips_init( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 0 app_config = self._build_app_config() default_file = make_upload_file(tenant_id=app_config.tenant_id) @@ -678,16 +663,14 @@ def _record_commit(session: Session) -> None: get_thread_messages_length = MagicMock(return_value=0) thread_data: dict[str, object] = {} - monkeypatch.setattr(generator, "_init_generate_records", init_records) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.get_thread_messages_length", get_thread_messages_length - ) + monkeypatch.setattr(generator._runtime.chat_records, "initialize", init_records) + monkeypatch.setattr(generator._runtime.chat_records, "dialogue_count", get_thread_messages_length) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.MessageBasedAppQueueManager", + "services.workflow.execution.adapters.chatflow.app_generator.MessageBasedAppQueueManager", lambda **kwargs: SimpleNamespace(**kwargs), ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.current_app", + "services.workflow.execution.adapters.chatflow.app_generator.current_app", SimpleNamespace(_get_current_object=lambda: SimpleNamespace(name="flask")), ) @@ -706,11 +689,7 @@ def join(self, 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=sqlite_engine, session=sqlite_session), - ) + monkeypatch.setattr("services.workflow.execution.adapters.chatflow.app_generator.threading.Thread", _Thread) monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory") def handle_response(**kwargs): @@ -721,7 +700,7 @@ def handle_response(**kwargs): monkeypatch.setattr(generator, "_handle_advanced_chat_response", handle_response) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppGenerateResponseConverter.convert", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppGenerateResponseConverter.convert", lambda response, invoke_from: response, ) @@ -730,7 +709,6 @@ def handle_response(**kwargs): user=_make_end_user(end_user_id="user"), invoke_from=InvokeFrom.WEB_APP, application_generate_entity=application_generate_entity, - session=sqlite_session, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), conversation=conversation, @@ -740,7 +718,7 @@ def handle_response(**kwargs): assert response == {"raw": True} init_records.assert_not_called() - get_thread_messages_length.assert_called_once_with(conversation.id, session=sqlite_session) + get_thread_messages_length.assert_called_once_with(conversation.id) assert thread_data["started"] is True assert thread_data["joined"] is True assert thread_data["join_timeout"] == 300 @@ -752,9 +730,16 @@ def handle_response(**kwargs): ) def test_generate_worker_raises_when_workflow_not_found( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 app_config = self._build_app_config() @@ -772,38 +757,50 @@ def test_generate_worker_raises_when_workflow_not_found( workflow_run_id="run-id", ) - generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) - generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) - @contextmanager def _fake_context[**P](*args: P.args, **kwargs: P.kwargs): yield - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), + "services.workflow.execution.adapters.chatflow.app_generator.preserve_flask_contexts", _fake_context ) - with pytest.raises(ValueError, match="Workflow not found"): - generator._generate_worker( - flask_app=SimpleNamespace(), - application_generate_entity=application_generate_entity, - queue_manager=MagicMock(), - conversation_id="conv", - message_id="msg", - context=SimpleNamespace(), - variable_loader=SimpleNamespace(), - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - graph_runtime_state=None, - ) + sqlite_session.add_all( + [ + _app_generator_make_app(), + _make_conversation(conversation_id="conv"), + _make_message(message_id="msg", conversation_id="conv"), + ] + ) + sqlite_session.commit() + queue = MagicMock() + generator._generate_worker( + system_user_id="session-id", + flask_app=SimpleNamespace(), + application_generate_entity=application_generate_entity, + queue_manager=queue, + conversation_id="conv", + message_id="msg", + context=SimpleNamespace(), + variable_loader=SimpleNamespace(), + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + graph_engine_layers=(), + graph_runtime_state=None, + ) + queue.publish_error.assert_called_once() def test_generate_worker_raises_when_app_not_found_for_internal_call( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 app_config = self._build_app_config() @@ -821,8 +818,6 @@ def test_generate_worker_raises_when_app_not_found_for_internal_call( workflow_run_id="run-id", ) - generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) - generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) sqlite_session.add(_app_generator_make_workflow()) sqlite_session.commit() @@ -830,31 +825,38 @@ def test_generate_worker_raises_when_app_not_found_for_internal_call( def _fake_context[**P](*args: P.args, **kwargs: P.kwargs): yield - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), + "services.workflow.execution.adapters.chatflow.app_generator.preserve_flask_contexts", _fake_context ) - with pytest.raises(ValueError, match="App not found"): - generator._generate_worker( - flask_app=SimpleNamespace(), - application_generate_entity=application_generate_entity, - queue_manager=MagicMock(), - conversation_id="conv", - message_id="msg", - context=SimpleNamespace(), - variable_loader=SimpleNamespace(), - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - graph_runtime_state=None, - ) + queue = MagicMock() + generator._generate_worker( + system_user_id="session-id", + flask_app=SimpleNamespace(), + application_generate_entity=application_generate_entity, + queue_manager=queue, + conversation_id="conv", + message_id="msg", + context=SimpleNamespace(), + variable_loader=SimpleNamespace(), + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + graph_engine_layers=(), + graph_runtime_state=None, + ) + queue.publish_error.assert_called_once() def test_generate_worker_handles_stopped_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 app_config = self._build_app_config() @@ -873,34 +875,40 @@ def test_generate_worker_handles_stopped_error( ) queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) - generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) @contextmanager def _fake_context[**P](*args: P.args, **kwargs: P.kwargs): yield - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) + monkeypatch.setattr( + "services.workflow.execution.adapters.chatflow.app_generator.preserve_flask_contexts", _fake_context + ) - sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow()]) + sqlite_session.add_all( + [ + _app_generator_make_app(), + _app_generator_make_workflow(), + _make_conversation(conversation_id="conv"), + _make_message(message_id="msg", conversation_id="conv"), + ] + ) sqlite_session.commit() class _Runner: def __init__(self, **kwargs): - _ = kwargs + assert kwargs["runtime"] is workflow_runtime def prepare(self): raise GenerateTaskStoppedError() - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", _Runner) - restore_workflow_run_graph = MagicMock() - monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppRunner", _Runner ) + restore_workflow_run_graph = MagicMock(side_effect=lambda workflow, _run_id: workflow) + monkeypatch.setattr(generator._runtime.contexts, "restore_graph", restore_workflow_run_graph) generator._generate_worker( + system_user_id="session-id", flask_app=SimpleNamespace(), application_generate_entity=application_generate_entity, queue_manager=queue_manager, @@ -910,17 +918,25 @@ def prepare(self): variable_loader=SimpleNamespace(), workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), + graph_engine_layers=(), graph_runtime_state=SimpleNamespace(), ) queue_manager.publish_error.assert_not_called() - assert restore_workflow_run_graph.call_args.kwargs["workflow"].id == "workflow-id" - assert restore_workflow_run_graph.call_args.kwargs["workflow_run_id"] == "run-id" + assert restore_workflow_run_graph.call_args.args[0].id == "workflow-id" + assert restore_workflow_run_graph.call_args.args[1] == "run-id" def test_generate_worker_handles_validation_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 app_config = self._build_app_config() @@ -949,16 +965,23 @@ class _ValidationModel(BaseModel): raise AssertionError("validation error should be created") queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) - generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) - sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow()]) + sqlite_session.add_all( + [ + _app_generator_make_app(), + _app_generator_make_workflow(), + _make_conversation(conversation_id="conv"), + _make_message(message_id="msg", conversation_id="conv"), + ] + ) sqlite_session.commit() @contextmanager def _fake_context[**P](*args: P.args, **kwargs: P.kwargs): yield - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) + monkeypatch.setattr( + "services.workflow.execution.adapters.chatflow.app_generator.preserve_flask_contexts", _fake_context + ) class _Runner: def __init__(self, **kwargs): @@ -967,13 +990,12 @@ def __init__(self, **kwargs): def prepare(self): raise validation_error - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", _Runner) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppRunner", _Runner ) generator._generate_worker( + system_user_id="session-id", flask_app=SimpleNamespace(), application_generate_entity=application_generate_entity, queue_manager=queue_manager, @@ -983,16 +1005,29 @@ def prepare(self): variable_loader=SimpleNamespace(), workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), + graph_engine_layers=(), graph_runtime_state=None, ) queue_manager.publish_error.assert_called_once() def test_generate_worker_handles_value_and_unknown_errors( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): app_config = self._build_app_config() - sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow()]) + sqlite_session.add_all( + [ + _app_generator_make_app(), + _app_generator_make_workflow(), + _make_conversation(conversation_id="conv"), + _make_message(message_id="msg", conversation_id="conv"), + ] + ) sqlite_session.commit() @contextmanager @@ -1010,7 +1045,9 @@ def prepare(self): return _Runner for raised_error in [ValueError("bad input"), RuntimeError("unexpected")]: - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( task_id="task", @@ -1027,21 +1064,18 @@ def prepare(self): ) queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) - generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", - _make_runner(raised_error), + "services.workflow.execution.adapters.chatflow.app_generator.preserve_flask_contexts", _fake_context ) - apply_config_overrides(monkeypatch, DEBUG=True) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppRunner", + _make_runner(raised_error), ) + apply_config_overrides(monkeypatch, DEBUG=True) generator._generate_worker( + system_user_id="session-id", flask_app=SimpleNamespace(), application_generate_entity=application_generate_entity, queue_manager=queue_manager, @@ -1051,13 +1085,18 @@ def prepare(self): variable_loader=SimpleNamespace(), workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), + graph_engine_layers=(), graph_runtime_state=None, ) queue_manager.publish_error.assert_called_once() - def test_handle_response_closed_file_raises_stopped(self, monkeypatch: pytest.MonkeyPatch): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_handle_response_closed_file_raises_stopped( + self, monkeypatch: pytest.MonkeyPatch, *, workflow_runtime: WorkflowRuntime + ): + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 app_config = WorkflowUIBasedAppConfig( @@ -1090,7 +1129,7 @@ def process(self): raise ValueError("I/O operation on closed file.") monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppGenerateTaskPipeline", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppGenerateTaskPipeline", _Pipeline, ) @@ -1113,9 +1152,15 @@ def process(self): ) def test_handle_response_re_raises_value_error( - self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + self, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 app_config = self._build_app_config() application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( @@ -1139,9 +1184,11 @@ def __init__(self, **kwargs): def process(self): raise ValueError("other error") - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.AdvancedChatAppGenerateTaskPipeline", _Pipeline) + monkeypatch.setattr( + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppGenerateTaskPipeline", _Pipeline + ) - with caplog.at_level(logging.ERROR, logger="core.app.apps.advanced_chat.app_generator"): + with caplog.at_level(logging.ERROR, logger="services.workflow.execution.adapters.chatflow.app_generator"): with pytest.raises(ValueError, match="other error"): generator._handle_advanced_chat_response( application_generate_entity=application_generate_entity, @@ -1163,9 +1210,16 @@ def process(self): assert "Failed to process generate task pipeline, conversation_id: conv" in caplog.messages def test_generate_worker_handles_invoke_auth_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) generator._dialogue_count = 1 app_config = WorkflowUIBasedAppConfig( @@ -1192,8 +1246,6 @@ def test_generate_worker_handles_invoke_auth_error( queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) - generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow(), _make_end_user()]) sqlite_session.commit() @@ -1206,20 +1258,20 @@ def prepare(self): raise InvokeAuthorizationError("bad key") - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", _Runner) + monkeypatch.setattr( + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppRunner", _Runner + ) @contextmanager def _fake_context[**P](*args: P.args, **kwargs: P.kwargs): yield - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), + "services.workflow.execution.adapters.chatflow.app_generator.preserve_flask_contexts", _fake_context ) generator._generate_worker( + system_user_id="session-id", flask_app=SimpleNamespace(), application_generate_entity=application_generate_entity, queue_manager=queue_manager, @@ -1229,15 +1281,23 @@ def _fake_context[**P](*args: P.args, **kwargs: P.kwargs): variable_loader=SimpleNamespace(), workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), + graph_engine_layers=(), graph_runtime_state=None, ) assert queue_manager.publish_error.called def test_generate_debugger_enables_retrieve_source( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", @@ -1249,11 +1309,11 @@ def test_generate_debugger_enables_retrieve_source( ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppConfigManager.get_app_config", lambda app_model, workflow: app_config, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", + "services.workflow.execution.adapters.chatflow.app_generator.FileUploadConfigManager.convert", lambda features_dict, is_vision=False: None, ) dummy_trace_queue_manager = type( @@ -1266,21 +1326,9 @@ def test_generate_debugger_enables_retrieve_source( }, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", + "services.workflow.execution.adapters.chatflow.app_generator.TraceQueueManager", dummy_trace_queue_manager, ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=unbound_session), - ) captured = {} @@ -1302,7 +1350,6 @@ def _fake_generate(**kwargs): invoke_from=InvokeFrom.DEBUGGER, workflow_run_id="run-id", streaming=False, - session=unbound_session, ) assert result == {"ok": True} @@ -1310,9 +1357,16 @@ def _fake_generate(**kwargs): assert captured["application_generate_entity"].query == "hello" def test_generate_service_api_sets_parent_message_id( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + sqlite_engine: _app_generator_Engine, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", @@ -1324,11 +1378,11 @@ def test_generate_service_api_sets_parent_message_id( ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppConfigManager.get_app_config", lambda app_model, workflow: app_config, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", + "services.workflow.execution.adapters.chatflow.app_generator.FileUploadConfigManager.convert", lambda features_dict, is_vision=False: None, ) dummy_trace_queue_manager = type( @@ -1341,21 +1395,9 @@ def test_generate_service_api_sets_parent_message_id( }, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", + "services.workflow.execution.adapters.chatflow.app_generator.TraceQueueManager", dummy_trace_queue_manager, ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=sqlite_engine, session=unbound_session), - ) captured = {} @@ -1377,7 +1419,6 @@ def _fake_generate(**kwargs): invoke_from=InvokeFrom.SERVICE_API, workflow_run_id="run-id", streaming=False, - session=unbound_session, ) assert captured["application_generate_entity"].parent_message_id == UUID_NIL @@ -1396,9 +1437,15 @@ def _build_app_config() -> WorkflowUIBasedAppConfig: ) def test_resume_restores_trace_manager_when_missing( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + *, + workflow_runtime: WorkflowRuntime, ): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( task_id="task", app_config=self._build_app_config(), @@ -1425,7 +1472,7 @@ def test_resume_restores_trace_manager_when_missing( }, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", + "services.workflow.execution.adapters.chatflow.app_generator.TraceQueueManager", dummy_trace_queue_manager, ) captured_entity: AdvancedChatAppGenerateEntity | None = None @@ -1443,7 +1490,6 @@ def _fake_generate(**kwargs): user=_make_end_user(), conversation=_make_conversation(), message=_make_message(), - session=unbound_session, application_generate_entity=application_generate_entity, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), @@ -1457,8 +1503,16 @@ def _fake_generate(**kwargs): assert trace_manager.app_id == "app-id" assert trace_manager.user_id == "session-id" - def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_resume_preserves_existing_trace_manager( + self, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + *, + workflow_runtime: WorkflowRuntime, + ): + generator = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) existing_trace_manager = SimpleNamespace(app_id="existing-app", user_id="existing-user") application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( task_id="task", @@ -1491,7 +1545,6 @@ def _fake_generate(**kwargs): user=_make_end_user(), conversation=_make_conversation(), message=_make_message(), - session=unbound_session, application_generate_entity=application_generate_entity, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), @@ -1565,24 +1618,31 @@ def _account() -> _advanced_chat_app_generator_Account: return account -def test_init_generate_records_sets_conversation_metadata(sqlite_session: Session): +def test_init_generate_records_sets_conversation_metadata( + sqlite_session: Session, *, workflow_runtime: WorkflowRuntime +): app_config = _make_app_config() entity = _make_generate_entity(app_config) - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) - - conversation, _ = generator._init_generate_records( - entity, - conversation=None, - session=sqlite_session, + sqlite_session.add(_app()) + sqlite_session.commit() + conversation, message = workflow_runtime.chat_records.initialize( + tenant_id=app_config.tenant_id, + app_id=app_config.app_id, + conversation_id=None, + seed=prepare_chat_records(entity, "{}"), ) - assert entity.conversation_id == conversation.id + assert message.conversation_id == conversation.id assert conversation.id is not None - assert entity.is_new_conversation is True + assert inspect(conversation).detached + assert inspect(message).detached + assert conversation.from_end_user_id == entity.user_id -def test_init_generate_records_marks_existing_conversation(sqlite_session: Session): +def test_init_generate_records_marks_existing_conversation( + sqlite_session: Session, *, workflow_runtime: WorkflowRuntime +): app_config = _make_app_config() entity = _make_generate_entity(app_config) @@ -1605,26 +1665,27 @@ def test_init_generate_records_marks_existing_conversation(sqlite_session: Sessi from_account_id=None, ) existing_conversation.id = "existing-conversation-id" - sqlite_session.add(existing_conversation) - sqlite_session.flush() - - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + sqlite_session.add_all([_app(), existing_conversation]) + sqlite_session.commit() - conversation, _ = generator._init_generate_records( - entity, - conversation=existing_conversation, - session=sqlite_session, + conversation, message = workflow_runtime.chat_records.initialize( + tenant_id=app_config.tenant_id, + app_id=app_config.app_id, + conversation_id=existing_conversation.id, + seed=prepare_chat_records(entity, "{}"), ) - assert entity.conversation_id == "existing-conversation-id" - assert conversation is existing_conversation - assert entity.is_new_conversation is False + assert message.conversation_id == "existing-conversation-id" + assert conversation.id == existing_conversation.id + assert conversation.name == "existing" def test_generate_falls_back_to_new_conversation_when_conversation_missing( monkeypatch: pytest.MonkeyPatch, - sqlite_engine: _advanced_chat_app_generator_Engine, + sqlite_engine: _app_generator_Engine, sqlite_session: Session, + *, + workflow_runtime: WorkflowRuntime, ): app_config = _make_app_config() workflow = _workflow() @@ -1635,34 +1696,25 @@ def raise_conversation_not_exists(**_kwargs): raise ConversationNotExistsError() monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.ConversationService.get_conversation", + workflow_runtime.chat_records, + "conversation", raise_conversation_not_exists, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", + "services.workflow.execution.adapters.chatflow.app_generator.FileUploadConfigManager.convert", lambda *_args, **_kwargs: None, ) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + "services.workflow.execution.adapters.chatflow.app_generator.AdvancedChatAppConfigManager.get_app_config", lambda **_kwargs: app_config, ) - monkeypatch.setattr(app_generator_module, "db", SimpleNamespace(engine=sqlite_engine)) trace_manager = object.__new__(TraceQueueManager) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", + "services.workflow.execution.adapters.chatflow.app_generator.TraceQueueManager", lambda **_kwargs: trace_manager, ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **_kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **_kwargs: SimpleNamespace(), - ) captured: dict[str, object] = {} - session = sqlite_session def fake_generate(self, **kwargs): captured.update(kwargs) @@ -1670,7 +1722,9 @@ def fake_generate(self, **kwargs): monkeypatch.setattr(AdvancedChatAppGenerator, "_generate", fake_generate) - result = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run).generate( + result = AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ).generate( app_model=app_model, workflow=workflow, user=user, @@ -1678,18 +1732,16 @@ def fake_generate(self, **kwargs): invoke_from=InvokeFrom.SERVICE_API, workflow_run_id="workflow-run-id", streaming=False, - session=session, ) assert result == {"status": "ok"} assert captured["conversation"] is None - assert captured["session"] is session application_generate_entity = captured["application_generate_entity"] assert isinstance(application_generate_entity, AdvancedChatAppGenerateEntity) assert application_generate_entity.conversation_id is None -def test_message_cycle_manager_uses_new_conversation_flag(monkeypatch: pytest.MonkeyPatch): +def test_message_cycle_manager_uses_new_conversation_flag(app_records, monkeypatch: pytest.MonkeyPatch): app_config = _make_app_config() entity = _make_generate_entity(app_config) entity.conversation_id = "existing-conversation-id" @@ -1713,7 +1765,7 @@ def fake_thread[**P](*args: P.args, **kwargs: P.kwargs): monkeypatch.setattr(message_cycle_manager, "Timer", fake_thread) - manager = MessageCycleManager(application_generate_entity=entity, task_state=MagicMock()) + manager = MessageCycleManager(records=app_records, application_generate_entity=entity, task_state=MagicMock()) thread = manager.generate_conversation_name(conversation_id="existing-conversation-id", query="hello") assert thread is captured["thread"] @@ -1994,7 +2046,7 @@ def _fake_generate(**kwargs): assert resumed_state.outputs == baseline_outputs -def test_advanced_chat_pause_resume_matches_baseline(mocker: MockerFixture, unbound_session: Session): +def test_advanced_chat_pause_resume_matches_baseline(mocker: MockerFixture, *, workflow_runtime: WorkflowRuntime): _patch_tool_node(mocker) baseline_state = _build_runtime_state("adv-baseline") @@ -2011,7 +2063,9 @@ def test_advanced_chat_pause_resume_matches_baseline(mocker: MockerFixture, unbo resumed_state = RuntimeState.from_snapshot(snapshot) - generator = adv_app_gen_module.AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = adv_app_gen_module.AdvancedChatAppGenerator( + runtime=workflow_runtime, execution_driver=workflow_runtime.execution_driver + ) def _fake_generate(**kwargs): state: RuntimeState = kwargs["graph_runtime_state"] @@ -2026,7 +2080,6 @@ def _fake_generate(**kwargs): user=_pause_resume_make_account(), conversation=_pause_resume_make_conversation(), message=_pause_resume_make_message(), - session=unbound_session, application_generate_entity=SimpleNamespace( stream=False, invoke_from=InvokeFrom.SERVICE_API, diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner.py index 0d3c132407046a..3837eae03ba15c 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner.py @@ -1,17 +1,13 @@ import json -from contextlib import contextmanager from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest -from sqlalchemy import event, select -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy import select +from sqlalchemy.orm import Session -import core.app.apps.advanced_chat.app_runner as module +import services.workflow.execution.adapters.chatflow.app_runner as module from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat import app_runner as app_runner_module -from core.app.apps.advanced_chat.app_runner import AdvancedChatAppRunner from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueStopEvent from core.moderation.base import ModerationError @@ -21,11 +17,16 @@ from models import Conversation as _app_runner_conversation_variables_Conversation from models import ConversationVariable from models import Message as _app_runner_conversation_variables_Message +from models.annotation_reply import AnnotationReply +from models.enums import ConversationFromSource from models.model import App as _app_runner_input_moderation_App -from models.model import AppMode, IconType, MessageAnnotation +from models.model import AppMode, IconType from models.model import Conversation as _app_runner_input_moderation_Conversation from models.model import Message as _app_runner_input_moderation_Message from models.workflow import Workflow, WorkflowType +from repositories.workflow.definition_repository import workflow_snapshot +from services.workflow.execution.adapters.chatflow.app_runner import AdvancedChatAppRunner +from services.workflow.execution.ports import WorkflowRuntime """SQLite-backed conversation-variable synchronization tests for AdvancedChatAppRunner.""" @@ -48,7 +49,7 @@ def _variable(variable_id: str, name: str, value: str): ) -def _runner(workflow_variables: list[object]) -> AdvancedChatAppRunner: +def _runner(workflow_variables: list[object], workflow_runtime: WorkflowRuntime) -> AdvancedChatAppRunner: workflow = Workflow( id="55555555-5555-5555-5555-555555555555", tenant_id="66666666-6666-6666-6666-666666666666", @@ -68,18 +69,10 @@ def _runner(workflow_variables: list[object]) -> AdvancedChatAppRunner: message=_app_runner_conversation_variables_Message(id="77777777-7777-7777-7777-777777777777"), dialogue_count=1, variable_loader=MagicMock(), - workflow=workflow, + workflow=workflow_snapshot(workflow), system_user_id="44444444-4444-4444-4444-444444444444", app=_app_runner_conversation_variables_App(id=APP_ID, tenant_id=workflow.tenant_id), - ) - - -def _bind_runner_sessions(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - engine = sqlite_session.get_bind() - monkeypatch.setattr( - app_runner_module, - "create_session", - lambda: Session(engine, expire_on_commit=False), + runtime=workflow_runtime, ) @@ -94,17 +87,34 @@ def _persist_variable(session: Session, *, variable: object, conversation_id: st session.commit() +def _persist_conversation(session: Session, conversation_id: str = CONVERSATION_ID) -> None: + session.add( + _app_runner_conversation_variables_Conversation( + id=conversation_id, + app_id=APP_ID, + mode=AppMode.ADVANCED_CHAT, + name="Conversation", + _inputs={}, + from_source=ConversationFromSource.API, + ) + ) + session.commit() + + class TestAppRunnerConversationVariablesCases: @staticmethod @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) - def test_missing_conversation_variables_are_added(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + def test_missing_conversation_variables_are_added( + sqlite_session: Session, *, workflow_runtime: WorkflowRuntime + ) -> None: existing_variable = _variable(VAR_1_ID, "existing_var", "default1") new_variable = _variable(VAR_2_ID, "new_var", "default2") + _persist_conversation(sqlite_session) + _persist_conversation(sqlite_session, OTHER_CONVERSATION_ID) _persist_variable(sqlite_session, variable=existing_variable) _persist_variable(sqlite_session, variable=new_variable, conversation_id=OTHER_CONVERSATION_ID) - _bind_runner_sessions(monkeypatch, sqlite_session) - variables = _runner([existing_variable, new_variable])._initialize_conversation_variables() + variables = _runner([existing_variable, new_variable], workflow_runtime)._initialize_conversation_variables() assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] persisted = sqlite_session.scalars( @@ -116,14 +126,14 @@ def test_missing_conversation_variables_are_added(monkeypatch: pytest.MonkeyPatc @staticmethod @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) - def test_no_variables_creates_all(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + def test_no_variables_creates_all(sqlite_session: Session, *, workflow_runtime: WorkflowRuntime) -> None: workflow_variables = [ _variable(VAR_1_ID, "var1", "default1"), _variable(VAR_2_ID, "var2", "default2"), ] - _bind_runner_sessions(monkeypatch, sqlite_session) + _persist_conversation(sqlite_session) - variables = _runner(workflow_variables)._initialize_conversation_variables() + variables = _runner(workflow_variables, workflow_runtime)._initialize_conversation_variables() assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] persisted = sqlite_session.scalars(select(ConversationVariable).order_by(ConversationVariable.id)).all() @@ -131,16 +141,16 @@ def test_no_variables_creates_all(monkeypatch: pytest.MonkeyPatch, sqlite_sessio @staticmethod @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) - def test_all_variables_exist_no_changes(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + def test_all_variables_exist_no_changes(sqlite_session: Session, *, workflow_runtime: WorkflowRuntime) -> None: workflow_variables = [ _variable(VAR_1_ID, "var1", "default1"), _variable(VAR_2_ID, "var2", "default2"), ] + _persist_conversation(sqlite_session) for variable in workflow_variables: _persist_variable(sqlite_session, variable=variable) - _bind_runner_sessions(monkeypatch, sqlite_session) - variables = _runner(workflow_variables)._initialize_conversation_variables() + variables = _runner(workflow_variables, workflow_runtime)._initialize_conversation_variables() assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] persisted = sqlite_session.scalars(select(ConversationVariable)).all() @@ -149,12 +159,12 @@ def test_all_variables_exist_no_changes(monkeypatch: pytest.MonkeyPatch, sqlite_ @staticmethod @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) def test_non_uuid_conversation_variable_is_created_once( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + sqlite_session: Session, *, workflow_runtime: WorkflowRuntime ) -> None: author_id = "opt-comp-prompt-var" variable = _variable(author_id, "optimization_comparison_prompt", "-") - _bind_runner_sessions(monkeypatch, sqlite_session) - runner = _runner([variable]) + _persist_conversation(sqlite_session) + runner = _runner([variable], workflow_runtime) first = runner._initialize_conversation_variables() second = runner._initialize_conversation_variables() @@ -185,7 +195,7 @@ def test_non_uuid_conversation_variable_is_created_once( def _patch_common_run_deps(runner: AdvancedChatAppRunner): """Context manager that patches common heavy deps used by run().""" return patch.multiple( - "core.app.apps.advanced_chat.app_runner", + "services.workflow.execution.adapters.chatflow.app_runner", RedisChannel=MagicMock(), redis_client=MagicMock(), WorkflowEntry=MagicMock(**{"return_value.run.return_value": iter([])}), @@ -196,7 +206,7 @@ def _patch_common_run_deps(runner: AdvancedChatAppRunner): class TestAppRunnerInputModerationCases: @staticmethod @pytest.fixture - def build_runner(sqlite_session: Session): + def build_runner(sqlite_session: Session, *, workflow_runtime: WorkflowRuntime): """Construct a minimal AdvancedChatAppRunner with heavy dependencies mocked.""" app_id = str(uuid4()) workflow_id = str(uuid4()) @@ -205,7 +215,14 @@ def build_runner(sqlite_session: Session): # Mocks for constructor args mock_queue_manager = MagicMock() - conversation = _app_runner_input_moderation_Conversation(id=str(uuid4()), app_id=app_id) + conversation = _app_runner_input_moderation_Conversation( + id=str(uuid4()), + app_id=app_id, + mode=AppMode.ADVANCED_CHAT, + name="Test", + _inputs={}, + from_source=ConversationFromSource.API, + ) message = _app_runner_input_moderation_Message(id=str(uuid4()), app_id=app_id, conversation_id=conversation.id) workflow = Workflow( id=workflow_id, @@ -233,7 +250,7 @@ def build_runner(sqlite_session: Session): enable_site=False, enable_api=False, ) - sqlite_session.add(app) + sqlite_session.add_all([app, conversation]) sqlite_session.commit() gen = AdvancedChatAppGenerateEntity( @@ -255,9 +272,10 @@ def build_runner(sqlite_session: Session): message=message, dialogue_count=1, variable_loader=MagicMock(), - workflow=workflow, + workflow=workflow_snapshot(workflow), system_user_id=str(uuid4()), app=app, + runtime=workflow_runtime, ) return runner @@ -304,7 +322,7 @@ def test_run_applies_overridden_inputs_and_query_from_moderation(build_runner): return_value=(True, overridden_inputs, overridden_query), ) as mock_moderate, patch.object(runner, "handle_annotation_reply", return_value=False) as mock_anno, - patch.object(runner, "_init_graph", return_value=MagicMock()) as mock_init_graph, + patch.object(runner._graphs, "build", return_value=MagicMock()) as mock_init_graph, ): runner.prepare() @@ -339,7 +357,7 @@ def test_run_returns_early_when_direct_output_via_handle_input_moderation(build_ runner.application_generate_entity.query, ), ) as mock_handle, - patch.object(runner, "_init_graph") as mock_init_graph, + patch.object(runner._graphs, "build") as mock_init_graph, patch.object(runner, "handle_annotation_reply") as mock_anno, ): runner.prepare() @@ -350,59 +368,39 @@ def test_run_returns_early_when_direct_output_via_handle_input_moderation(build_ mock_init_graph.assert_not_called() @staticmethod - def test_run_publishes_annotation_after_commit(build_runner, sqlite_engine: Engine): + def test_run_publishes_annotation_after_commit(build_runner): runner = build_runner events: list[str] = [] - def record_commit(session: Session) -> None: - if session.get_bind() is sqlite_engine: - events.append("commit") - - event.listen(Session, "after_commit", record_commit) - annotation_reply = MessageAnnotation( - app_id=runner._app.id, - question="question", - content="annotated answer", - account_id=str(uuid4()), - ) + def query_annotation(**_kwargs): + events.append("commit") + return AnnotationReply(id="annotation-id", content="annotated answer") def publish(event): if isinstance(event, QueueAnnotationReplyEvent): events.append("publish") with ( - _patch_common_run_deps(runner), - patch.object( - runner, - "handle_input_moderation", - return_value=( - False, - runner.application_generate_entity.inputs, - runner.application_generate_entity.query, - ), - ), - patch.object(runner, "handle_annotation_reply", return_value=annotation_reply), - patch.object(runner, "_publish_event", side_effect=publish), + patch.object(runner, "moderation_for_inputs", return_value=(False, {}, "raw-query")), + patch.object(runner._runtime.annotation_replies, "query", side_effect=query_annotation), + patch.object(runner._events, "_publish_event", side_effect=publish), patch.object(runner, "_complete_with_stream_output"), ): - runner.prepare() - event.remove(Session, "after_commit", record_commit) + assert runner.prepare() is None assert events == ["commit", "publish"] @staticmethod - def test_run_closes_scoped_session_before_workflow_run(build_runner, sqlite_session_factory: sessionmaker[Session]): + def test_run_closes_scoped_session_before_workflow_run(build_runner): runner = build_runner events = [] - @contextmanager - def observed_session(): - with sqlite_session_factory() as session: - yield session - events.append("close") - workflow_entry = MagicMock() + def initialize_conversation_variables(): + events.append("close") + return [] + def run_workflow(**kwargs): events.append("run") return iter([]) @@ -410,7 +408,6 @@ def run_workflow(**kwargs): workflow_entry.run.side_effect = run_workflow with ( - patch.object(module, "create_session", observed_session), patch.object(module, "RedisChannel"), patch.object(module, "redis_client"), patch.object(module, "WorkflowEntry", return_value=workflow_entry), @@ -423,9 +420,9 @@ def run_workflow(**kwargs): runner.application_generate_entity.query, ), ), - patch.object(runner, "handle_annotation_reply", return_value=False), - patch.object(runner, "_initialize_conversation_variables", return_value=[]), - patch.object(runner, "_init_graph", return_value=MagicMock()), + patch.object(runner, "handle_annotation_reply", return_value=None), + patch.object(runner, "_initialize_conversation_variables", side_effect=initialize_conversation_variables), + patch.object(runner._graphs, "build", return_value=MagicMock()), ): prepared = runner.prepare() assert prepared is not None diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py index d8aca97e83eb0f..3b4f9a0099297a 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py @@ -1,23 +1,17 @@ from __future__ import annotations -from contextlib import nullcontext -from datetime import datetime +from dataclasses import dataclass, replace +from datetime import timedelta from types import SimpleNamespace from unittest import mock import pytest -from pytest_mock import MockerFixture -from sqlalchemy import select +from sqlalchemy import create_engine, event, select from sqlalchemy.orm import Session, sessionmaker from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat import generate_task_pipeline as pipeline_module -from core.app.apps.advanced_chat.generate_task_pipeline import ( - AdvancedChatAppGenerateTaskPipeline, - ConversationSnapshot, - MessageSnapshot, - WorkflowSnapshot, -) +from core.app.apps.draft_variable_saver import NoopDraftVariableSaver +from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom from core.app.entities.queue_entities import ( QueueAdvancedChatMessageEndEvent, @@ -64,30 +58,170 @@ from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.runtime import RuntimeState, VariablePool from libs.datetime_utils import naive_utc_now +from models.base import TypeBase from models.enums import CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom from models.execution_extra_content import HumanInputContent +from models.human_input import HumanInputForm, HumanInputFormRecipient from models.human_input_entities import UserActionConfig -from models.model import AppMode, EndUser, Message, MessageFile -from models.workflow import WorkflowRun +from models.model import App, AppMode, Conversation, EndUser, Message, MessageFile +from models.workflow import WorkflowNodeExecutionModel, WorkflowPause, WorkflowRun from models.workflow import WorkflowType as ModelWorkflowType +from repositories.app.generation_repository import AppGenerationRepository from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository +from repositories.workflow.runtime_context_repository import WorkflowRuntimeContextRepository +from services.app.generation.ports import ChatRecordSeed, ConversationSnapshot, MessageSnapshot, WorkflowSnapshot +from services.errors.message import MessageNotExistsError +from services.workflow.execution.adapters.chatflow import generate_task_pipeline as pipeline_module +from services.workflow.execution.adapters.chatflow.generate_task_pipeline import AdvancedChatAppGenerateTaskPipeline +from tests.helpers.unit.model_factories import make_app, make_conversation, make_end_user from tests.workflow_test_utils import build_test_variable_pool -def _build_generate_entity() -> AdvancedChatAppGenerateEntity: - return AdvancedChatAppGenerateEntity( - task_id="task-1", - app_config=WorkflowUIBasedAppConfig( - app_id="app-1", tenant_id="tenant-1", app_mode=AppMode.ADVANCED_CHAT, workflow_id="workflow-1" +@dataclass +class MessageCase: + pipeline: pipeline_module.AdvancedChatAppGenerateTaskPipeline + sessions: sessionmaker[Session] + + def message(self) -> Message: + with self.sessions() as session: + return session.get(Message, "message-1") + + +@pytest.fixture +def message_store(): + # Keep the injected message database separate from the globally configured test database. + engine = create_engine("sqlite://") + TypeBase.metadata.create_all( + engine, + tables=[ + model.__table__ + for model in ( + App, + Conversation, + Message, + MessageFile, + HumanInputForm, + HumanInputFormRecipient, + HumanInputContent, + WorkflowRun, + WorkflowPause, + WorkflowNodeExecutionModel, + ) + ], + ) + try: + yield sessionmaker(bind=engine, expire_on_commit=False) + finally: + engine.dispose() + + +@pytest.fixture +def app_records(message_store): + return AppGenerationRepository(message_store) + + +@pytest.fixture +def runtime_app_records(sqlite_session_factory: sessionmaker[Session]) -> AppGenerationRepository: + return AppGenerationRepository(sqlite_session_factory) + + +@pytest.fixture +def message_case(app_records, message_store, *, tool_providers) -> MessageCase: + with message_store.begin() as session: + session.add(make_app(mode=AppMode.ADVANCED_CHAT)) + session.add( + HumanInputForm( + id="form-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + node_id="node-1", + form_definition="{}", + rendered_content="Approve?", + expiration_time=naive_utc_now() + timedelta(days=1), + ) + ) + conversation, message = app_records.initialize( + tenant_id="tenant-1", + app_id="app-1", + conversation_id=None, + seed=ChatRecordSeed( + conversation={ + "id": "conversation-1", + "mode": AppMode.ADVANCED_CHAT, + "name": "New conversation", + "_inputs": {}, + "from_source": "api", + "from_end_user_id": "user-1", + "from_account_id": None, + }, + message={ + "id": "message-1", + "_inputs": {}, + "query": "hello", + "message": "", + "answer": "", + "message_unit_price": 0, + "answer_unit_price": 0, + "currency": "USD", + "from_source": "api", + "from_end_user_id": "user-1", + "invoke_from": InvokeFrom.WEB_APP, + }, + files=[], ), + ) + config = WorkflowUIBasedAppConfig( + tenant_id="tenant-1", + app_id="app-1", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow-1", + variables=[], + additional_features=AppAdditionalFeatures(), + ) + entity = AdvancedChatAppGenerateEntity( + task_id="task-1", + app_config=config, inputs={}, + query="hello", files=[], user_id="user-1", stream=True, invoke_from=InvokeFrom.WEB_APP, - query="hello", workflow_run_id="run-1", + extras={"trace_session_id": "trace-1"}, + ) + queue = MessageBasedAppQueueManager( + task_id=entity.task_id, + user_id=entity.user_id, + invoke_from=entity.invoke_from, + conversation_id=conversation.id, + app_mode=AppMode.ADVANCED_CHAT, + message_id=message.id, ) + pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( + contexts=WorkflowRuntimeContextRepository(message_store), + application_generate_entity=entity, + workflow=WorkflowSnapshot("workflow-1", "tenant-1", {}), + queue_manager=queue, + conversation=ConversationSnapshot.from_conversation(conversation), + message=MessageSnapshot.from_message(message), + user=make_end_user(end_user_id="user-1"), + stream=True, + dialogue_count=app_records.dialogue_count(conversation.id), + draft_var_saver_factory=lambda *_args, **_kwargs: NoopDraftVariableSaver(), + chat_records=app_records, + tool_providers=tool_providers, + ) + pipeline._graph_runtime_state = RuntimeState( + workflow_id="workflow-1", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-1") + ), + start_at=0, + ) + pipeline._workflow_run_id = "run-1" + return MessageCase(pipeline, message_store) def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline: @@ -100,381 +234,163 @@ def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline: return pipeline -def test_process_passes_message_id_to_conversation_name_generation() -> None: - pipeline = _build_pipeline() - pipeline._conversation_id = "conversation-1" - pipeline._application_generate_entity = SimpleNamespace(query="hello", trace_manager=None) - pipeline._message_cycle_manager = mock.Mock() - pipeline._base_task_pipeline = SimpleNamespace(stream=True) - pipeline._wrapper_process_stream_response = mock.Mock(return_value=iter(())) - pipeline._to_stream_response = mock.Mock(return_value="stream") - - result = pipeline.process() - - assert result == "stream" - pipeline._message_cycle_manager.generate_conversation_name.assert_called_once_with( - conversation_id="conversation-1", query="hello", message_id="message-1" +def test_process_passes_message_id_to_conversation_name_generation(message_case, monkeypatch): + pipeline = message_case.pipeline + calls = [] + monkeypatch.setattr( + pipeline._message_cycle_manager, "generate_conversation_name", lambda **kwargs: calls.append(kwargs) ) + monkeypatch.setattr(pipeline, "_wrapper_process_stream_response", lambda **kwargs: iter(())) + assert list(pipeline.process()) == [] + assert calls == [{"conversation_id": "conversation-1", "query": "hello", "message_id": "message-1"}] -def _message(*, status: MessageStatus, answer: str = "") -> Message: - return Message( - id="message-1", - status=status, - answer=answer, - invoke_from=InvokeFrom.WEB_APP, - from_end_user_id="user-1", - ) - - -def test_persist_human_input_extra_content_adds_record(sqlite_session: Session) -> None: - pipeline = _build_pipeline() - - pipeline._persist_human_input_extra_content(form_id="form-1") - - content = sqlite_session.scalar(select(HumanInputContent)) - assert content is not None - assert content.workflow_run_id == "run-1" - assert content.message_id == "message-1" - assert content.form_id == "form-1" - - -def test_persist_human_input_extra_content_skips_when_existing(sqlite_session: Session) -> None: - pipeline = _build_pipeline() - existing = HumanInputContent.new(workflow_run_id="run-1", message_id="message-1", form_id="form-1") - sqlite_session.add(existing) - sqlite_session.commit() - - pipeline._persist_human_input_extra_content(form_id="form-1") - - contents = list(sqlite_session.scalars(select(HumanInputContent))) - assert [content.id for content in contents] == [existing.id] +@pytest.mark.parametrize("form_exists", [True, False]) +def test_human_input_association_is_persisted_once(message_case, form_exists): + pipeline = message_case.pipeline + node = "node-1" if form_exists else "missing" + form_id = "form-1" if form_exists else None + pipeline._persist_human_input_extra_content(form_id=form_id, node_id=node) + pipeline._persist_human_input_extra_content(form_id=form_id, node_id=node) + with message_case.sessions() as session: + rows = list(session.scalars(select(HumanInputContent))) + assert [(row.message_id, row.form_id) for row in rows] == ([("message-1", "form-1")] if form_exists else []) @pytest.mark.parametrize("stop_after_pause", [False, True]) -def test_handle_workflow_paused_event_does_not_overwrite_a_concurrent_stop( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - stop_after_pause: bool, - mocker: MockerFixture, -) -> None: - run = WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - task_id="task-1", - type=ModelWorkflowType(WorkflowType.CHAT), - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - status=WorkflowExecutionStatus.PAUSED, - created_by_role=CreatorUserRole.END_USER, - created_by="user-1", - ) - sqlite_session.add(run) - sqlite_session.commit() - pipeline = _build_pipeline() - pipeline._application_generate_entity = _build_generate_entity() - pipeline._workflow_response_converter = mock.Mock() - pipeline._workflow_response_converter.workflow_pause_to_stream_response.return_value = ["pause"] - pipeline._ensure_graph_runtime_initialized = mock.Mock( - return_value=SimpleNamespace( - total_tokens=0, - node_run_steps=0, - ), - ) - pipeline._save_message = mock.Mock() - message = _message(status=MessageStatus.NORMAL) - pipeline._get_message = mock.Mock(return_value=message) - persist_content = mocker.patch.object(pipeline, "_persist_human_input_extra_content") - pipeline._base_task_pipeline = mock.Mock() - pipeline._base_task_pipeline.queue_manager = mock.Mock() - pipeline._message_saved_on_pause = False - +def test_pause_does_not_overwrite_a_concurrent_stop(message_case, stop_after_pause): + pipeline = message_case.pipeline + pipeline._workflow_response_converter.workflow_pause_to_stream_response = lambda **_kwargs: ["pause"] + with message_case.sessions.begin() as session: + session.add( + WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + task_id="task-1", + type=ModelWorkflowType(WorkflowType.CHAT), + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + status=WorkflowExecutionStatus.PAUSED, + created_by_role=CreatorUserRole.END_USER, + created_by="user-1", + ) + ) + list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) + pipeline._task_state.answer = "before" reason = HumanInputRequired( form_id="form-1", - form_content="content", + form_content="Approve?", inputs=[], actions=[], node_id="node-1", node_title="Approval", resolved_default_values={}, ) - event = QueueWorkflowPausedEvent(reasons=[reason], outputs={}, paused_nodes=["node-1"]) - - events = pipeline._handle_workflow_paused_event(event) - assert next(events) == "pause" + responses = pipeline._handle_workflow_paused_event( + QueueWorkflowPausedEvent(reasons=[reason], outputs={}, paused_nodes=["node-1"]) + ) + assert next(responses) == "pause" if stop_after_pause: - repository = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) + repository = DifyAPISQLAlchemyWorkflowRunRepository(message_case.sessions) repository.stop_paused_workflow_task(tenant_id="tenant-1", app_id="app-1", task_id="task-1", owner=None) - assert list(events) == [] - - persist_content.assert_called_once_with(form_id="form-1") - assert message.status == (MessageStatus.NORMAL if stop_after_pause else MessageStatus.PAUSED) - - -def test_resume_appends_chunks_to_paused_answer(unbound_session: Session) -> None: - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=None, - extras={}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conversation-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="message-1", - created_at=datetime(2024, 1, 1), - query="hello", - answer="before", - status=MessageStatus.PAUSED, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) - - stored_message = _message(status=MessageStatus.PAUSED, answer="before") - pipeline._get_message = mock.Mock(return_value=stored_message) - pipeline._recorded_files = [] - + assert list(responses) == [] + assert message_case.message().status == (MessageStatus.NORMAL if stop_after_pause else MessageStatus.PAUSED) + assert message_case.message().answer == "before" + with message_case.sessions() as session: + assert session.scalar(select(HumanInputContent.form_id)) == "form-1" + + +def test_resume_appends_chunks_and_resets_paused_status(message_case): + pipeline = message_case.pipeline + pipeline._task_state.answer = "before" + pipeline._save_message(paused=True) + pipeline._seed_task_state_from_message(MessageSnapshot.from_message(message_case.message())) list(pipeline._handle_text_chunk_event(QueueTextChunkEvent(text="after"))) - pipeline._save_message(session=unbound_session) - - assert stored_message.answer == "beforeafter" - assert stored_message.status == MessageStatus.NORMAL - - -def test_emit_message_trace_enqueues_message_trace() -> None: - """When trace_manager is provided, _emit_message_trace must enqueue a MESSAGE_TRACE task.""" - from core.ops.entities.trace_entity import TraceTaskName - - mock_trace_manager = mock.MagicMock() - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=mock_trace_manager, - extras={"trace_session_id": "session-trace-1"}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conv-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="msg-1", - created_at=datetime(2024, 1, 1), - query="hello", - answer="", - status=MessageStatus.NORMAL, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) - - stored_message = SimpleNamespace( - id="msg-1", - answer="", - status=MessageStatus.NORMAL, - updated_at=None, - provider_response_latency=0, - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - total_price=0, - currency="USD", - message_metadata=None, - invoke_from=InvokeFrom.WEB_APP, - from_account_id=None, - from_end_user_id="user-1", - ) - pipeline._get_message = mock.Mock(return_value=stored_message) - pipeline._recorded_files = [] - - pipeline._emit_message_trace() - - mock_trace_manager.add_trace_task.assert_called_once() - trace_task = mock_trace_manager.add_trace_task.call_args[0][0] - assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE - assert trace_task.conversation_id == "conv-1" - assert trace_task.message_id == "msg-1" - assert trace_task.kwargs.get("trace_session_id") == "session-trace-1" - - -def test_save_message_does_not_emit_trace() -> None: - """_save_message must not enqueue a trace; emitting is the terminal caller's responsibility. - - This keeps intermediate saves (e.g. on pause) from duplicating message, token, and TTFT metrics - when the message is later resumed. - """ - mock_trace_manager = mock.MagicMock() - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=mock_trace_manager, - extras={"trace_session_id": "session-trace-1"}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conv-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="msg-1", - created_at=datetime(2024, 1, 1), - query="hello", - answer="", - status=MessageStatus.NORMAL, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) - - stored_message = SimpleNamespace( - id="msg-1", - answer="", - status=MessageStatus.NORMAL, - updated_at=None, - provider_response_latency=0, - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - total_price=0, - currency="USD", - message_metadata=None, - invoke_from=InvokeFrom.WEB_APP, - from_account_id=None, - from_end_user_id="user-1", - ) - pipeline._get_message = mock.Mock(return_value=stored_message) - pipeline._recorded_files = [] - - pipeline._save_message(session=mock.Mock()) - - mock_trace_manager.add_trace_task.assert_not_called() - - -def test_emit_message_trace_skips_when_no_trace_manager() -> None: - """When trace_manager is None, _emit_message_trace and _save_message must not fail.""" - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=None, - extras={}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conv-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="msg-1", - created_at=datetime(2024, 1, 1), - query="hello", - answer="", - status=MessageStatus.NORMAL, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) - - stored_message = SimpleNamespace( - id="msg-1", - answer="", - status=MessageStatus.NORMAL, - updated_at=None, - provider_response_latency=0, - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - total_price=0, - currency="USD", - message_metadata=None, - invoke_from=InvokeFrom.WEB_APP, - from_account_id=None, - from_end_user_id="user-1", - ) - pipeline._get_message = mock.Mock(return_value=stored_message) - pipeline._recorded_files = [] - - # Should not raise - pipeline._save_message(session=mock.Mock()) + pipeline._save_message() + assert message_case.message().answer == "beforeafter" + assert message_case.message().status == MessageStatus.NORMAL + + +def test_start_completion_files_and_error_use_initialization_database(message_case): + pipeline = message_case.pipeline + list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) + assert message_case.message().workflow_run_id == "run-1" + pipeline._task_state.answer = "done" + pipeline._recorded_files = [ + {"type": "image", "transfer_method": "remote_url", "remote_url": "https://example.com/a.png"} + ] + responses = list(pipeline._handle_workflow_succeeded_event(QueueWorkflowSucceededEvent(outputs={}))) + assert responses[0].event == StreamEvent.MESSAGE_END + assert message_case.message().answer == "done" + with message_case.sessions() as session: + file = session.scalar(select(MessageFile)) + assert file.message_id == "message-1" + assert file.created_by == "user-1" + list(pipeline._handle_error_event(QueueErrorEvent(error=ValueError("failed")))) + assert message_case.message().status == MessageStatus.ERROR + assert message_case.message().error == "failed" + + +def test_failed_attachment_write_rolls_back_answer_and_status(message_case): + pipeline = message_case.pipeline + pipeline._task_state.answer = "uncommitted" + pipeline._recorded_files = [ + {"type": "image", "transfer_method": "remote_url", "remote_url": "https://example.com/a.png"} + ] + engine = message_case.sessions.kw["bind"] + + def fail_file(_conn, _cursor, sql, _params, _ctx, _many): + if sql.startswith("INSERT INTO message_files"): + raise RuntimeError("file write failed") + + event.listen(engine, "before_cursor_execute", fail_file) + try: + with pytest.raises(RuntimeError, match="file write failed"): + pipeline._save_message(paused=True) + finally: + event.remove(engine, "before_cursor_execute", fail_file) + assert message_case.message().answer == "" + assert message_case.message().status == MessageStatus.NORMAL + + +def test_deleted_message_does_not_suppress_start_or_error_event(message_case): + with message_case.sessions.begin() as session: + session.delete(session.get(Message, "message-1")) + pipeline = message_case.pipeline + assert list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) + responses = list(pipeline._handle_error_event(QueueErrorEvent(error=ValueError("failed")))) + assert len(responses) == 1 + assert str(responses[0].err) == "failed" + + +@pytest.mark.parametrize("field", ["tenant_id", "app_id", "conversation_id"]) +def test_message_writes_check_owner_chain(message_case, field): + pipeline = message_case.pipeline + pipeline._message_identity = replace(pipeline._message_identity, **{field: "foreign"}) + with pytest.raises(MessageNotExistsError): + pipeline._save_message() + assert message_case.message().answer == "" + + +def test_trace_is_enqueued_after_message_commit(message_case): + pipeline = message_case.pipeline + traces = [] + + class TraceRecorder: + def add_trace_task(self, task): + assert message_case.message().answer == "done" + traces.append(task) + + pipeline._application_generate_entity.trace_manager = TraceRecorder() + pipeline._task_state.answer = "done" + pipeline._save_message() + assert traces == [] pipeline._emit_message_trace() + assert len(traces) == 1 + assert traces[0].message_id == "message-1" + assert traces[0].kwargs["trace_session_id"] == "trace-1" def test_workflow_succeeded_emits_message_end_before_workflow_finished() -> None: @@ -571,7 +487,7 @@ def test_process_stream_response_breaks_after_workflow_partial_success() -> None pipeline._base_task_pipeline.ping_stream_response.assert_not_called() -def _make_pipeline(): +def _make_pipeline(app_records=None): app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", app_id="app", @@ -606,6 +522,8 @@ def _make_pipeline(): user = EndUser(tenant_id="tenant", type="session", name="tester", session_id="session") pipeline = AdvancedChatAppGenerateTaskPipeline( + contexts=mock.MagicMock(), + chat_records=app_records or mock.MagicMock(), application_generate_entity=application_generate_entity, workflow=workflow, queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), @@ -615,6 +533,7 @@ def _make_pipeline(): stream=False, dialogue_count=1, draft_var_saver_factory=lambda **kwargs: None, + tool_providers=mock.MagicMock(), ) return pipeline @@ -649,6 +568,10 @@ def _persist_message(session: Session, *, message_id: str = "message-id") -> Mes status=MessageStatus.PAUSED, ) message.id = message_id + if session.get(App, "app") is None: + session.add(make_app(app_id="app", tenant_id="tenant")) + if session.get(Conversation, "conv-id") is None: + session.add(make_conversation(conversation_id="conv-id", app_id="app", inputs={}, from_source="api")) session.add(message) session.commit() return message @@ -805,8 +728,8 @@ def test_handle_error_event(self): assert isinstance(responses[0], ValueError) - def test_handle_workflow_started_event_sets_run_id(self, sqlite_session: Session): - pipeline = _make_pipeline() + def test_handle_workflow_started_event_sets_run_id(self, runtime_app_records, sqlite_session: Session): + pipeline = _make_pipeline(runtime_app_records) message = _persist_message(sqlite_session) other_message = _persist_message(sqlite_session, message_id="other-message-id") pipeline._graph_runtime_state = RuntimeState( @@ -1092,14 +1015,14 @@ def append_new_token(self, text): assert result is False assert seen == ["token"] - def test_handle_retriever_and_annotation_events(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): + def test_handle_retriever_and_annotation_events(self, monkeypatch: pytest.MonkeyPatch): pipeline = _make_pipeline() calls = {"retriever": 0, "annotation": 0} def _hit_retriever(event): calls["retriever"] += 1 - def _hit_annotation(_manager, event, session): + def _hit_annotation(_manager, event): calls["annotation"] += 1 pipeline._message_cycle_manager.handle_retriever_resources = _hit_retriever @@ -1108,8 +1031,6 @@ def _hit_annotation(_manager, event, session): retriever_event = QueueRetrieverResourcesEvent(retriever_resources=[]) annotation_event = QueueAnnotationReplyEvent(message_annotation_id="ann") - monkeypatch.setattr(pipeline, "_database_session", lambda: nullcontext(unbound_session)) - assert list(pipeline._handle_retriever_resources_event(retriever_event)) == [] assert list(pipeline._handle_annotation_reply_event(annotation_event)) == [] assert calls == {"retriever": 1, "annotation": 1} @@ -1153,8 +1074,8 @@ def test_handle_human_input_events(self): assert list(pipeline._handle_human_input_form_timeout_event(timeout_event)) == ["timeout"] assert persisted == ["saved"] - def test_save_message_preserves_full_answer_and_sets_usage(self, sqlite_session: Session): - pipeline = _make_pipeline() + def test_save_message_preserves_full_answer_and_sets_usage(self, runtime_app_records, sqlite_session: Session): + pipeline = _make_pipeline(runtime_app_records) pipeline._recorded_files = [ { "type": FileType.IMAGE, @@ -1179,8 +1100,8 @@ def test_save_message_preserves_full_answer_and_sets_usage(self, sqlite_session: start_at=0.0, ) - pipeline._save_message(session=sqlite_session, graph_runtime_state=graph_runtime_state) - sqlite_session.commit() + pipeline._save_message(graph_runtime_state=graph_runtime_state) + sqlite_session.refresh(message) assert message.status == MessageStatus.NORMAL assert message.answer == "![img](http://example.com/file.png) hello ![inline](http://llm.com/img.jpg)" diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py index d7bd1f2d32f07f..9807bb3806bff4 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py @@ -1,78 +1,49 @@ +"""Unit tests for AgentAppGenerator.generate() and its worker thread. + +Mirrors the agent_chat generator tests: every collaborator (config manager, +model converter, queue manager, thread, response converter, the agent backend +client stack) is patched at the module level, the generate entity class is +patched so no real pydantic entity is built, and the worker's flask-context +manager is replaced with a no-op so the thread body can run inline. +""" + from __future__ import annotations import contextlib import inspect from decimal import Decimal from types import SimpleNamespace -from typing import Any -from unittest.mock import MagicMock +from unittest.mock import create_autospec import pytest from pytest_mock import MockerFixture -from sqlalchemy import event, select from sqlalchemy.orm import Session -import core.app.apps.agent_app.app_generator as module -import core.app.features.annotation_reply.annotation_reply as annotation_mod -import core.moderation.input_moderation as input_moderation_mod -from core.app.apps.agent_app import app_generator -from core.app.apps.agent_app.app_generator import ( - AgentAppGenerator, - AgentAppGeneratorError, - AgentAppNotPublishedError, -) from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.entities.queue_entities import ( - QueueAnnotationReplyEvent, - QueueLLMChunkEvent, - QueueMessageEndEvent, -) -from core.moderation.base import ModerationError -from models import Account as _app_generator_Account -from models import AppModelConfig -from models.account import Account as _resolve_agent_Account -from models.agent import ( - Agent, - AgentConfigDraft, - AgentConfigDraftType, - AgentConfigSnapshot, - AgentConfigVersionKind, - AgentScope, - AgentSource, - AgentStatus, - AgentWorkingResourceStatus, - AgentWorkspaceBinding, -) +from core.app.entities.queue_entities import QueueAnnotationReplyEvent +from models import Account, AppModelConfig +from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource, AgentStatus from models.agent_config_entities import AgentSoulConfig from models.enums import AppStatus, ConversationFromSource from models.model import App, AppMode, Conversation, Message, MessageAnnotation -from services.agent.workspace_service import ( - AgentWorkspaceBindingGenerationMismatchError, - AgentWorkspaceService, +from models.tool_runtime_contracts import WorkflowToolQueries +from services.app.generation.adapters.agent_app import ( + AgentAppGenerator, + AgentAppGeneratorError, ) -from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_message - -"""Unit tests for AgentAppGenerator.generate() and its worker thread. - -Mirrors the agent_chat generator tests: every collaborator (config manager, -model converter, queue manager, thread, response converter, the agent backend -client stack) is patched at the module level, the generate entity class is -patched so no real pydantic entity is built, and the worker's flask-context -manager is replaced with a no-op so the thread body can run inline. -""" +from services.app.generation.agent_config import AgentAppConfiguration, AgentAppConfigurations +MODULE = "services.app.generation.adapters.agent_app" -MODULE = "core.app.apps.agent_app.app_generator" - -def _app_generator_account(user_id: str = "user") -> _app_generator_Account: - account = _app_generator_Account(name="User", email=f"{user_id}@example.com") +def _account(user_id: str = "user") -> Account: + account = Account(name="User", email=f"{user_id}@example.com") account.id = user_id return account -def _app_generator_app(*, app_model_config_id: str | None = None) -> App: +def _app(*, app_model_config_id: str | None = None) -> App: return App( id="app1", tenant_id="tenant", @@ -100,7 +71,7 @@ def _agent(*, agent_id: str = "agent1") -> Agent: ) -def _app_generator_snapshot(*, snapshot_id: str = "snap1", agent_id: str = "agent1") -> AgentConfigSnapshot: +def _snapshot(*, snapshot_id: str = "snap1", agent_id: str = "agent1") -> AgentConfigSnapshot: return AgentConfigSnapshot( id=snapshot_id, tenant_id="tenant", @@ -147,1036 +118,357 @@ def _message(*, query: str = "query") -> Message: _CURRENT_SESSION: Session | None = None +@pytest.fixture(autouse=True) +def _bind_real_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): + global _CURRENT_SESSION + _CURRENT_SESSION = sqlite_session + sqlite_session.add(_app()) + sqlite_session.commit() + yield + _CURRENT_SESSION = None + + def _session() -> Session: assert _CURRENT_SESSION is not None return _CURRENT_SESSION -class TestAppGeneratorCases: - @staticmethod - @pytest.fixture(autouse=True) - def _bind_real_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): - global _CURRENT_SESSION - _CURRENT_SESSION = sqlite_session - monkeypatch.setattr(module.db, "session", sqlite_session) - yield - _CURRENT_SESSION = None - - @staticmethod - @pytest.fixture - def generator(mocker: MockerFixture) -> AgentAppGenerator: - gen = AgentAppGenerator() - mocker.patch(f"{MODULE}.current_app", new=mocker.MagicMock(_get_current_object=mocker.MagicMock())) - mocker.patch(f"{MODULE}.contextvars.copy_context", return_value="ctx") - return gen - - class TestGenerateGuards: - def test_rejects_blocking_mode(self, generator: AgentAppGenerator): - with pytest.raises(AgentAppGeneratorError, match="only supports streaming"): - generator.generate( - app_model=_app_generator_app(), - user=_app_generator_account("u"), - args={}, - invoke_from=InvokeFrom.WEB_APP, - session=_session(), - streaming=False, - ) - - def test_requires_query(self, generator: AgentAppGenerator): - with pytest.raises(AgentAppGeneratorError, match="query is required"): - generator.generate( - app_model=_app_generator_app(), - user=_app_generator_account("u"), - args={"inputs": {}}, - invoke_from=InvokeFrom.WEB_APP, - session=_session(), - ) - - def test_rejects_blank_query(self, generator: AgentAppGenerator): - with pytest.raises(AgentAppGeneratorError, match="query is required"): - generator.generate( - app_model=_app_generator_app(), - user=_app_generator_account("u"), - args={"query": " ", "inputs": {}}, - invoke_from=InvokeFrom.WEB_APP, - session=_session(), - ) - - class TestGenerateSuccess: - def test_session_scope_config_version_id_preserves_draft_or_snapshot_id(self): - assert ( - AgentAppGenerator._session_scope_config_version_id( - invoke_from=InvokeFrom.DEBUGGER, config_version_id="draft-1" - ) - == "draft-1" - ) - assert ( - AgentAppGenerator._session_scope_config_version_id( - invoke_from=InvokeFrom.WEB_APP, config_version_id="snapshot-1" - ) - == "snapshot-1" - ) - - def test_generate_orchestrates_and_starts_worker(self, generator, mocker: MockerFixture): - config = AppModelConfig(app_id="app1") - config.id = "config-1" - session = _session() - session.add(config) - session.commit() - app_model = _app_generator_app(app_model_config_id=config.id) - user = _app_generator_account() - - generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) - generator._prepare_user_inputs = mocker.MagicMock(return_value={"x": 1}) - generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) - generator._handle_response = mocker.MagicMock(return_value="raw-response") +@pytest.fixture +def generator(tool_providers, mocker: MockerFixture, *, annotation_replies, app_records) -> AgentAppGenerator: + gen = AgentAppGenerator( + agent_configs=create_autospec(AgentAppConfigurations, instance=True, spec_set=True), + annotations=annotation_replies, + records=app_records, + tool_providers=tool_providers, + workflow_queries=create_autospec(WorkflowToolQueries, instance=True, spec_set=True), + ) + mocker.patch(f"{MODULE}.current_app", new=mocker.MagicMock(_get_current_object=mocker.MagicMock())) + mocker.patch(f"{MODULE}.contextvars.copy_context", return_value="ctx") + gen._agent_configs.resolve.return_value = AgentAppConfiguration( + "agent1", "snap1", "snapshot", AgentSoulConfig(), "home-1" + ) + gen._agent_configs.version.return_value = AgentAppConfiguration("a", "s", "snapshot", AgentSoulConfig(), "home-1") + return gen - mocker.patch( - f"{MODULE}.AgentAppConfigManager.get_app_config", - return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), - ) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) - file_upload_config = mocker.MagicMock() - mocker.patch(f"{MODULE}.FileUploadConfigManager.convert", return_value=file_upload_config) - parsed_file = mocker.MagicMock() - build_files = mocker.patch(f"{MODULE}.file_factory.build_from_mappings", return_value=[parsed_file]) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - generate_entity = mocker.patch( - f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") - ) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - thread_obj = mocker.MagicMock() - thread_constructor = mocker.patch(f"{MODULE}.threading.Thread", return_value=thread_obj) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - file_mappings = [ - { - "type": "image", - "transfer_method": "local_file", - "url": "", - "upload_file_id": "upload-file-1", - } - ] - - result = generator.generate( - app_model=app_model, - user=user, - args={"query": "hello", "inputs": {"name": "world"}, "files": file_mappings}, - invoke_from=InvokeFrom.WEB_APP, - session=session, - streaming=True, - ) - assert result == {"result": "ok"} - thread_obj.start.assert_called_once() - worker_call = thread_constructor.call_args - inspect.signature(worker_call.kwargs["target"]).bind(**worker_call.kwargs["kwargs"]) - generator._resolve_agent.assert_called_once_with( - app_model, +class TestGenerateGuards: + def test_rejects_blocking_mode(self, generator: AgentAppGenerator): + with pytest.raises(AgentAppGeneratorError, match="only supports streaming"): + generator.generate( + app_model=_app(), + user=_account("u"), + args={}, invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=user, - session=session, - conversation=None, - ) - assert session.get(AppModelConfig, "config-1") is config - build_files.assert_called_once() - assert build_files.call_args.kwargs["mappings"] == file_mappings - assert generate_entity.call_args.kwargs["files"] == [parsed_file] - assert generate_entity.call_args.kwargs["file_upload_config"] is file_upload_config - assert "agent_runtime_exit_intent" not in generate_entity.call_args.kwargs - - def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): - app_model = _app_generator_app() - generator._resolve_agent = mocker.MagicMock( - return_value=(_agent(agent_id="a"), "snap1", "snapshot", AgentSoulConfig()) + streaming=False, ) - generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) - generator._handle_response = mocker.MagicMock(return_value="raw") - get_conv = mocker.patch(f"{MODULE}.ConversationService.get_conversation", return_value=_conversation()) - mocker.patch(f"{MODULE}.AgentAppConfigManager.get_app_config", return_value=mocker.MagicMock(variables=[])) - mocker.patch(f"{MODULE}.load_annotation_reply_config", return_value={"enabled": False}) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - session = _session() - user = _app_generator_account() + def test_requires_query(self, generator: AgentAppGenerator): + with pytest.raises(AgentAppGeneratorError, match="query is required"): generator.generate( - app_model=app_model, - user=user, - args={"query": "hi", "inputs": {}, "conversation_id": "conv"}, + app_model=_app(), + user=_account("u"), + args={"inputs": {}}, invoke_from=InvokeFrom.WEB_APP, - session=session, - streaming=True, - ) - - get_conv.assert_called_once_with( - app_model=app_model, - conversation_id="conv", - user=user, - session=session, - ) - assert generator._resolve_agent.call_args.kwargs["conversation"].id == "conv" - assert generator._init_generate_records.call_args.kwargs["session"] is session - - def test_generate_does_not_include_trace_session_id_in_extras( - self, generator: AgentAppGenerator, mocker: MockerFixture - ): - app_model = _app_generator_app() - user = _app_generator_account() - - generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) - generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) - generator._handle_response = mocker.MagicMock(return_value="raw-response") - - mocker.patch( - f"{MODULE}.AgentAppConfigManager.get_app_config", - return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), ) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - generate_entity = mocker.patch( - f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") - ) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + def test_rejects_blank_query(self, generator: AgentAppGenerator): + with pytest.raises(AgentAppGeneratorError, match="query is required"): generator.generate( - app_model=app_model, - user=user, - args={"query": "hello", "inputs": {}, "trace_session_id": "session-1"}, + app_model=_app(), + user=_account("u"), + args={"query": " ", "inputs": {}}, invoke_from=InvokeFrom.WEB_APP, - session=_session(), - streaming=True, - ) - - assert generate_entity.call_args.kwargs["extras"] == {"auto_generate_conversation_name": True} - - class TestGenerateWorker: - @pytest.fixture(autouse=True) - def patch_context(self, mocker: MockerFixture): - @contextlib.contextmanager - def ctx_manager[**P](*args: P.args, **kwargs: P.kwargs): - yield - - mocker.patch("libs.flask_utils.preserve_flask_contexts", ctx_manager) - - def _wire( - self, - generator: AgentAppGenerator, - mocker: MockerFixture, - *, - run_side_effect=None, - handled=False, - guard_query="query", - ): - generator._get_conversation = mocker.MagicMock(return_value=_conversation()) - generator._get_message = mocker.MagicMock(return_value=_message()) - generator._run_input_guards = mocker.MagicMock(return_value=(handled, guard_query, None)) - resolved_agent = _agent(agent_id="a") - resolved_config = _app_generator_snapshot(snapshot_id="s", agent_id="a") - resolver_sessions: list[Session] = [] - - def resolve_agent_by_id(**kwargs): - resolver_sessions.append(kwargs["session"]) - return resolved_agent, resolved_config, AgentSoulConfig() - - generator._resolve_agent_by_id = mocker.MagicMock(side_effect=resolve_agent_by_id) - session = _session() - if session.get(App, "app1") is None: - session.add(_app_generator_app()) - session.commit() - mocker.patch(f"{MODULE}.DifyRunContext", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppRuntimeRequestBuilder", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.create_agent_backend_run_client", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentBackendRunEventAdapter", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppWorkspaceStore", return_value=mocker.MagicMock()) - runner = mocker.MagicMock() - if run_side_effect is not None: - runner.run.side_effect = run_side_effect - mocker.patch(f"{MODULE}.AgentAppRunner", return_value=runner) - return runner, resolver_sessions - - def _call( - self, - generator, - mocker: MockerFixture, - queue_manager, - *, - query="query", - session_scope_config_version_id="s", - files=(), - file_upload_config=None, - ): - generator._generate_worker( - flask_app=mocker.MagicMock(), - context=mocker.MagicMock(), - application_generate_entity=mocker.MagicMock( - app_config=SimpleNamespace(app_id="app1", tenant_id="tenant"), - agent_id="a", - agent_config_snapshot_id="s", - agent_session_scope_config_version_id=session_scope_config_version_id, - model_conf=mocker.MagicMock(model="m"), - query=query, - files=files, - file_upload_config=file_upload_config, - ), - queue_manager=queue_manager, - conversation_id="conv", - message_id="msg", - user_from=UserFrom.END_USER, - ) - - def test_happy_path_runs_backend(self, generator: AgentAppGenerator, mocker: MockerFixture): - runner, resolver_sessions = self._wire(generator, mocker) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - runner.run.assert_called_once() - assert resolver_sessions == [generator._resolve_agent_by_id.call_args.kwargs["session"]] - assert resolver_sessions[0].get_bind() is not None - assert runner.run.call_args.kwargs["home_snapshot_id"] == "home-1" - assert "home_snapshot_ref" not in runner.run.call_args.kwargs - queue_manager.publish_error.assert_not_called() - - def test_worker_passes_session_scope_config_version_to_runner(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker) - queue_manager = mocker.MagicMock() - - self._call(generator, mocker, queue_manager, session_scope_config_version_id=None) - - assert runner.run.call_args.kwargs["agent_config_snapshot_id"] == "s" - assert runner.run.call_args.kwargs["session_scope_snapshot_id"] is None - - def test_worker_passes_files_to_backend_runner_without_rewriting_query(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker, guard_query="你看得见这张图片吗") - queue_manager = mocker.MagicMock() - files = [mocker.MagicMock(), mocker.MagicMock()] - - self._call( - generator, - mocker, - queue_manager, - query="你看得见这张图片吗", - files=files, - ) - - assert runner.run.call_args.kwargs["query"] == "你看得见这张图片吗" - assert runner.run.call_args.kwargs["files"] == files - assert runner.run.call_args.kwargs["image_detail_config"] is None - - def test_input_guard_short_circuit_skips_backend(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker, handled=True) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - runner.run.assert_not_called() - - def test_annotation_reply_publishes_after_guard_transaction_commits(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker, handled=True) - annotation_reply = MessageAnnotation( - app_id="app1", - question="query", - content="annotated answer", - account_id="user", - ) - generator._run_input_guards.return_value = (True, "query", annotation_reply) - events: list[str] = [] - queue_manager = mocker.MagicMock() - - def record_commit(_session: Session) -> None: - events.append("commit") - - def publish(event, *_args): - if isinstance(event, QueueAnnotationReplyEvent): - events.append("publish") - - queue_manager.publish.side_effect = publish - - event.listen(type(_session()), "after_commit", record_commit) - try: - self._call(generator, mocker, queue_manager) - finally: - event.remove(type(_session()), "after_commit", record_commit) - - assert events == ["commit", "publish"] - runner.run.assert_not_called() - - def test_generate_task_stopped_is_swallowed(self, generator, mocker: MockerFixture): - self._wire(generator, mocker, run_side_effect=GenerateTaskStoppedError()) - # A locally stopped consumer does not prove runtime cancellation. The - # runner must hand off the typed terminal it actually observed. - mocker.patch( - "core.ops.base_trace_instance.application_tracing_hooks", - side_effect=AssertionError("stop intent must not invent a tracing terminal"), - ) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - queue_manager.publish_error.assert_not_called() - - def test_unexpected_error_is_published(self, generator: AgentAppGenerator, mocker: MockerFixture): - self._wire(generator, mocker, run_side_effect=ValueError("boom")) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - assert queue_manager.publish_error.called - - -"""Unit tests for AgentAppGenerator._run_input_guards. - -The guards apply content moderation + annotation reply before the Agent backend -call, mirroring the EasyUI chat runners. They can short-circuit the turn (a -blocked/preset moderation answer or a matched annotation) and publish a direct -answer, or pass through a possibly moderation-sanitized query. -""" - - -class _FakeQueueManager: - def __init__(self) -> None: - self.events: list[Any] = [] - - def publish(self, event: Any, _from: Any) -> None: - self.events.append(event) - - -def _make_entity(query: str = "hello") -> SimpleNamespace: - return SimpleNamespace( - app_config=SimpleNamespace(app_id="app-1", tenant_id="tenant-1"), - model_conf=SimpleNamespace(model="gpt-4o-mini"), - inputs={}, - query=query, - trace_manager=None, - user_id="user-1", - invoke_from=SimpleNamespace(), - ) - - -def _input_guards_app() -> App: - return make_app(name="Agent App", mode=AppMode.AGENT_CHAT, icon_type=None) - - -def _input_guards_message() -> Message: - return make_message(message_id="msg-1") - - -def _patch_moderation(monkeypatch: pytest.MonkeyPatch, *, returns=None, raises: Exception | None = None) -> None: - class _FakeModeration: - def check(self, **kwargs: Any): - if raises is not None: - raise raises - return returns - - monkeypatch.setattr(input_moderation_mod, "InputModeration", _FakeModeration) - - -def _patch_annotation(monkeypatch: pytest.MonkeyPatch, *, reply=None) -> None: - class _FakeAnnotation: - def query(self, **kwargs: Any): - return reply - - monkeypatch.setattr(annotation_mod, "AnnotationReplyFeature", _FakeAnnotation) - - -def _answer_text(events: list[Any]) -> str: - end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) - return end.llm_result.message.content - - -def _saved_user_query(events: list[Any]) -> str: - end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) - prompt_messages = end.llm_result.prompt_messages - assert len(prompt_messages) == 1 - return prompt_messages[0].content - - -class TestInputGuardsCases: - class TestRunInputGuards: - def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, returns=(False, {}, "hello")) - _patch_annotation(monkeypatch, reply=None) - qm = _FakeQueueManager() - - handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("hello"), - app_model=_input_guards_app(), - message=_input_guards_message(), - queue_manager=qm, - ) - - assert handled is False - assert query == "hello" - assert annotation_reply is None - assert qm.events == [] - - def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, returns=(True, {}, "[redacted]")) - _patch_annotation(monkeypatch, reply=None) - qm = _FakeQueueManager() - - handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("leak my secret"), - app_model=_input_guards_app(), - message=_input_guards_message(), - queue_manager=qm, - ) - - assert handled is False - assert query == "[redacted]" - assert annotation_reply is None - assert qm.events == [] - - def test_moderation_block_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, raises=ModerationError("blocked preset answer")) - _patch_annotation(monkeypatch, reply=None) - qm = _FakeQueueManager() - - handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("forbidden"), - app_model=_input_guards_app(), - message=_input_guards_message(), - queue_manager=qm, - ) - - assert handled is True - assert annotation_reply is None - assert any(isinstance(e, QueueLLMChunkEvent) for e in qm.events) - assert _answer_text(qm.events) == "blocked preset answer" - assert _saved_user_query(qm.events) == "forbidden" - - def test_annotation_hit_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, returns=(False, {}, "what is your name")) - annotation = MessageAnnotation( - app_id="app-1", - question="what is your name", - content="I am the annotated Iris.", - account_id="account-1", ) - annotation.id = "anno-1" - _patch_annotation(monkeypatch, reply=annotation) - qm = _FakeQueueManager() - - handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("what is your name"), - app_model=_input_guards_app(), - message=_input_guards_message(), - queue_manager=qm, - ) - - assert handled is True - assert annotation_reply.id == "anno-1" - assert qm.events == [] - - -"""Unit tests for AgentAppGenerator agent/snapshot resolution. - -Covers the DB-backed resolution helpers (the bound roster Agent + its published -Agent Soul snapshot) including every not-found error path, using persisted rows -in the shared SQLite test database. -""" - -_SOUL_DICT = { - "model": { - "plugin_id": "langgenius/openai", - "model_provider": "langgenius/openai/openai", - "model": "gpt-4o-mini", - }, - "prompt": {"system_prompt": "You are Iris."}, -} - - -def _resolve_agent_agent( - *, - agent_id: str = "agent-1", - scope: AgentScope = AgentScope.ROSTER, - source: AgentSource = AgentSource.AGENT_APP, - active_snapshot_id: str | None = "snap-1", - published: bool = True, -) -> Agent: - return Agent( - id=agent_id, - tenant_id="t1", - name=f"Agent {agent_id}", - scope=scope, - source=source, - app_id="app-1" if scope == AgentScope.ROSTER else None, - backing_app_id="app-1" if scope == AgentScope.WORKFLOW_ONLY else None, - active_config_snapshot_id=active_snapshot_id, - active_config_has_model=True, - active_config_is_published=published, - created_by="creator-1", - updated_by="updater-1", - ) - - -def _resolve_agent_snapshot( - *, snapshot_id: str = "snap-1", agent_id: str = "agent-1", home_snapshot_id: str = "home-1" -) -> AgentConfigSnapshot: - return AgentConfigSnapshot( - id=snapshot_id, - tenant_id="t1", - agent_id=agent_id, - version=1, - config_snapshot=AgentSoulConfig.model_validate(_SOUL_DICT), - home_snapshot_id=home_snapshot_id, - created_by="creator-1", - ) +class TestGenerateSuccess: + def test_generate_orchestrates_and_starts_worker(self, generator, mocker: MockerFixture): + config = AppModelConfig(app_id="app1") + config.id = "config-1" + session = _session() + session.add(config) + session.commit() + app_model = _app(app_model_config_id=config.id) + user = _account() -def _resolve_agent_app() -> App: - return make_app( - tenant_id="t1", - name="Agent App", - mode=AppMode.AGENT_CHAT, - icon_type=None, - enable_site=False, - enable_api=False, - ) - - -def _resolve_agent_account() -> _resolve_agent_Account: - return make_account(account_id="user-1", name="Agent User", email="agent-user@example.com") - - -def _resolve_agent_conversation(*, binding_id: str) -> Conversation: - return make_conversation(agent_workspace_binding_id=binding_id) + generator._prepare_user_inputs = mocker.MagicMock(return_value={"x": 1}) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) + generator._handle_response = mocker.MagicMock(return_value="raw-response") + mocker.patch( + f"{MODULE}.AgentAppConfigManager.get_app_config", + return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), + ) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) + file_upload_config = mocker.MagicMock() + mocker.patch(f"{MODULE}.FileUploadConfigManager.convert", return_value=file_upload_config) + parsed_file = mocker.MagicMock() + build_files = mocker.patch(f"{MODULE}.file_factory.build_from_mappings", return_value=[parsed_file]) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + generate_entity = mocker.patch( + f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") + ) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + thread_obj = mocker.MagicMock() + thread_constructor = mocker.patch(f"{MODULE}.threading.Thread", return_value=thread_obj) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + file_mappings = [ + { + "type": "image", + "transfer_method": "local_file", + "url": "", + "upload_file_id": "upload-file-1", + } + ] + + result = generator.generate( + app_model=app_model, + user=user, + args={"query": "hello", "inputs": {"name": "world"}, "files": file_mappings}, + invoke_from=InvokeFrom.WEB_APP, + streaming=True, + ) -def _binding( - *, - home_snapshot_id: str = "home-1", - version_id: str = "snap-1", - version_kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT, -) -> AgentWorkspaceBinding: - return AgentWorkspaceBinding( - id="binding-1", - tenant_id="t1", - app_id="app-1", - workspace_id="workspace-1", - agent_id="agent-1", - base_home_snapshot_id=home_snapshot_id, - agent_config_version_id=version_id, - agent_config_version_kind=version_kind, - backend_binding_ref="backend-binding-1", - status=AgentWorkingResourceStatus.ACTIVE, - ) + assert result == {"result": "ok"} + thread_obj.start.assert_called_once() + worker_call = thread_constructor.call_args + inspect.signature(worker_call.kwargs["target"]).bind(**worker_call.kwargs["kwargs"]) + generator._agent_configs.resolve.assert_called_once_with( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + account_id=user.id, + debug=False, + draft_type=None, + conversation_id=None, + ) + assert generate_entity.call_args.kwargs["agent_session_scope_config_version_id"] == "snap1" + assert session.get(AppModelConfig, "config-1").id == config.id + build_files.assert_called_once() + assert build_files.call_args.kwargs["mappings"] == file_mappings + assert generate_entity.call_args.kwargs["files"] == [parsed_file] + assert generate_entity.call_args.kwargs["file_upload_config"] is file_upload_config + assert "agent_runtime_exit_intent" not in generate_entity.call_args.kwargs + + def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): + app_model = _app() + generator._prepare_user_inputs = mocker.MagicMock(return_value={}) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) + generator._handle_response = mocker.MagicMock(return_value="raw") + get_conv = mocker.patch.object(generator._records, "conversation", return_value=_conversation()) + mocker.patch(f"{MODULE}.AgentAppConfigManager.get_app_config", return_value=mocker.MagicMock(variables=[])) + mocker.patch.object(generator._records, "annotation_config", return_value={"enabled": False}) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + user = _account() + + generator.generate( + app_model=app_model, + user=user, + args={"query": "hi", "inputs": {}, "conversation_id": "conv"}, + invoke_from=InvokeFrom.WEB_APP, + streaming=True, + ) + get_conv.assert_called_once_with( + app_id="app1", + conversation_id="conv", + account_id="user", + end_user_id=None, + ) + assert generator._agent_configs.resolve.call_args.kwargs["conversation_id"] == "conv" + assert "session" not in generator._init_generate_records.call_args.kwargs + + def test_generate_does_not_include_trace_session_id_in_extras( + self, generator: AgentAppGenerator, mocker: MockerFixture + ): + app_model = _app() + user = _account() + + generator._prepare_user_inputs = mocker.MagicMock(return_value={}) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) + generator._handle_response = mocker.MagicMock(return_value="raw-response") + + mocker.patch( + f"{MODULE}.AgentAppConfigManager.get_app_config", + return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), + ) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + generate_entity = mocker.patch( + f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") + ) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + + generator.generate( + app_model=app_model, + user=user, + args={"query": "hello", "inputs": {}, "trace_session_id": "session-1"}, + invoke_from=InvokeFrom.WEB_APP, + streaming=True, + ) -def _persist(session: Session, *models: object) -> None: - session.add_all(models) - session.commit() + assert generate_entity.call_args.kwargs["extras"] == {"auto_generate_conversation_name": True} -class TestResolveAgentCases: - class TestResolveAgentById: - def test_success_returns_agent_snapshot_soul(self, sqlite_session: Session): - agent = _resolve_agent_agent() - snapshot = _resolve_agent_snapshot() - sqlite_session.add_all([agent, snapshot]) - sqlite_session.commit() +class TestGenerateWorker: + @pytest.fixture(autouse=True) + def patch_context(self, mocker: MockerFixture): + @contextlib.contextmanager + def ctx_manager[**P](*args: P.args, **kwargs: P.kwargs): + yield + + mocker.patch("libs.flask_utils.preserve_flask_contexts", ctx_manager) + + def _wire( + self, + generator: AgentAppGenerator, + mocker: MockerFixture, + *, + run_side_effect=None, + handled=False, + guard_query="query", + ): + mocker.patch.object(generator._records, "load", return_value=(_app(), _conversation(), _message())) + generator._run_input_guards = mocker.MagicMock(return_value=(handled, guard_query, None)) + session = _session() + if session.get(App, "app1") is None: + session.add(_app()) + session.commit() + mocker.patch(f"{MODULE}.DifyRunContext", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppRuntimeRequestBuilder", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.create_agent_backend_run_client", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentBackendRunEventAdapter", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppWorkspaceStore", return_value=mocker.MagicMock()) + runner = mocker.MagicMock() + if run_side_effect is not None: + runner.run.side_effect = run_side_effect + mocker.patch(f"{MODULE}.AgentAppRunner", return_value=runner) + return runner, generator._agent_configs.version + + def _call( + self, + generator, + mocker: MockerFixture, + queue_manager, + *, + query="query", + session_scope_config_version_id="s", + files=(), + file_upload_config=None, + ): + generator._generate_worker( + flask_app=mocker.MagicMock(), + context=mocker.MagicMock(), + application_generate_entity=mocker.MagicMock( + app_config=SimpleNamespace(app_id="app1", tenant_id="tenant"), + agent_id="a", + agent_config_snapshot_id="s", + agent_config_version_kind="snapshot", + agent_session_scope_config_version_id=session_scope_config_version_id, + model_conf=mocker.MagicMock(model="m"), + query=query, + files=files, + file_upload_config=file_upload_config, + ), + queue_manager=queue_manager, + conversation_id="conv", + message_id="msg", + user_from=UserFrom.END_USER, + ) - resolved_agent, resolved_snapshot, soul = AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", agent_id="agent-1", snapshot_id="snap-1", session=sqlite_session - ) + def test_happy_path_runs_backend(self, generator: AgentAppGenerator, mocker: MockerFixture): + runner, resolver_sessions = self._wire(generator, mocker) + queue_manager = mocker.MagicMock() + self._call(generator, mocker, queue_manager) + runner.run.assert_called_once() + resolver_sessions.assert_called_once_with( + tenant_id="tenant", app_id="app1", agent_id="a", version_id="s", version_kind="snapshot", account_id=None + ) + assert runner.run.call_args.kwargs["home_snapshot_id"] == "home-1" + assert "home_snapshot_ref" not in runner.run.call_args.kwargs + queue_manager.publish_error.assert_not_called() - assert resolved_agent is agent - assert resolved_snapshot is snapshot - assert soul.prompt.system_prompt == "You are Iris." - assert soul.model is not None - assert soul.model.model == "gpt-4o-mini" - - def test_agent_missing_raises(self, sqlite_session: Session): - with pytest.raises(AgentAppGeneratorError, match="Agent not found"): - AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", agent_id="x", snapshot_id="snap-1", session=sqlite_session - ) - - def test_no_published_version_raises(self, sqlite_session: Session): - sqlite_session.add(_resolve_agent_agent(active_snapshot_id=None)) - sqlite_session.commit() - with pytest.raises(AgentAppGeneratorError, match="no published version"): - AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", agent_id="agent-1", snapshot_id=None, session=sqlite_session - ) - - def test_snapshot_missing_raises(self, sqlite_session: Session): - sqlite_session.add(_resolve_agent_agent()) - sqlite_session.commit() - with pytest.raises(AgentAppGeneratorError, match="published version not found"): - AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", - agent_id="agent-1", - snapshot_id="snap-1", - session=sqlite_session, - ) - - class TestResolveDebugDraft: - def test_missing_shared_draft_is_created_with_supplied_session(self, sqlite_session: Session): - agent = _resolve_agent_agent() - sqlite_session.add_all([agent, _resolve_agent_snapshot()]) - sqlite_session.commit() - flush_count = 0 - - def _record_flush(_session: Session, _context: object) -> None: - nonlocal flush_count - flush_count += 1 - - event.listen(sqlite_session, "after_flush", _record_flush) - - draft = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=None, - draft_id=None, - account_id=None, - session=sqlite_session, - ) + def test_worker_passes_session_scope_config_version_to_runner(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker) + queue_manager = mocker.MagicMock() - assert draft.draft_type == AgentConfigDraftType.DRAFT - assert draft.base_snapshot_id == "snap-1" - assert sqlite_session.scalar(select(AgentConfigDraft).where(AgentConfigDraft.id == draft.id)) is draft - assert flush_count == 1 - - def test_stale_workflow_only_shared_draft_is_rebased_to_active_snapshot(self, sqlite_session: Session): - agent = _resolve_agent_agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") - draft = AgentConfigDraft( - id="draft-1", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DRAFT, - account_id=None, - draft_owner_key="", - base_snapshot_id="snap-1", - home_snapshot_id="home-1", - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "old"}}), - ) - active_snapshot = AgentConfigSnapshot( - id="snap-2", - tenant_id="t1", - agent_id="agent-1", - version=2, - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "new"}}), - home_snapshot_id="home-2", - created_by="creator-1", - ) - sqlite_session.add_all([agent, draft, active_snapshot]) - sqlite_session.commit() - flush_count = 0 - - def _record_flush(_session: Session, _context: object) -> None: - nonlocal flush_count - flush_count += 1 - - event.listen(sqlite_session, "after_flush", _record_flush) - - resolved = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=None, - account_id=None, - session=sqlite_session, - ) + self._call(generator, mocker, queue_manager, session_scope_config_version_id=None) - assert resolved is draft - assert resolved.id == "draft-1" - assert resolved.base_snapshot_id == "snap-2" - assert resolved.home_snapshot_id == "home-2" - assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "new" - assert flush_count == 1 - - def test_build_draft_is_not_rebased_to_active_snapshot(self, sqlite_session: Session): - agent = _resolve_agent_agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") - draft = AgentConfigDraft( - id="build-draft-1", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DEBUG_BUILD, - account_id="account-1", - draft_owner_key="account-1", - base_snapshot_id="snap-1", - home_snapshot_id="home-build", - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "build edit"}}), - ) - sqlite_session.add_all([agent, draft]) - sqlite_session.commit() - - resolved = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=AgentConfigDraftType.DEBUG_BUILD.value, - account_id="account-1", - session=sqlite_session, - ) + assert runner.run.call_args.kwargs["agent_config_snapshot_id"] == "s" + assert runner.run.call_args.kwargs["session_scope_snapshot_id"] is None - assert resolved is draft - assert resolved.base_snapshot_id == "snap-1" - assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "build edit" - assert not sqlite_session.dirty - - def test_build_draft_uses_exact_draft_id(self, sqlite_session: Session): - agent = _resolve_agent_agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") - draft = AgentConfigDraft( - id="exact-build-draft", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DEBUG_BUILD, - account_id="account-1", - draft_owner_key="account-1", - base_snapshot_id="snap-1", - home_snapshot_id="home-build", - config_snapshot=AgentSoulConfig(), - ) - newer_decoy = AgentConfigDraft( - id="newer-build-draft", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DEBUG_BUILD, - account_id="account-1", - draft_owner_key="account-1-decoy", - base_snapshot_id="snap-2", - home_snapshot_id="home-decoy", - config_snapshot=AgentSoulConfig(prompt={"system_prompt": "decoy"}), - ) - sqlite_session.add_all([agent, draft, newer_decoy]) - sqlite_session.commit() - - resolved = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=AgentConfigDraftType.DEBUG_BUILD.value, - draft_id="exact-build-draft", - account_id="account-1", - session=sqlite_session, - ) + def test_worker_passes_files_to_backend_runner_without_rewriting_query(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker, guard_query="你看得见这张图片吗") + queue_manager = mocker.MagicMock() + files = [mocker.MagicMock(), mocker.MagicMock()] - assert resolved is draft - assert resolved.id == "exact-build-draft" + self._call( + generator, + mocker, + queue_manager, + query="你看得见这张图片吗", + files=files, + ) - class TestResolveAgent: - @pytest.fixture(autouse=True) - def _publish_visibility(self, monkeypatch: pytest.MonkeyPatch) -> None: - def is_publish_visible(*, agent: Agent, **_kwargs: object) -> bool: - return agent.active_config_is_published + assert runner.run.call_args.kwargs["query"] == "你看得见这张图片吗" + assert runner.run.call_args.kwargs["files"] == files + assert runner.run.call_args.kwargs["image_detail_config"] is None + + def test_input_guard_short_circuit_skips_backend(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker, handled=True) + queue_manager = mocker.MagicMock() + self._call(generator, mocker, queue_manager) + runner.run.assert_not_called() + + def test_annotation_reply_publishes_after_guard_result(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker, handled=True) + annotation_reply = MessageAnnotation( + app_id="app1", + question="query", + content="annotated answer", + account_id="user", + ) + generator._run_input_guards.return_value = (True, "query", annotation_reply) + events: list[str] = [] + queue_manager = mocker.MagicMock() - monkeypatch.setattr( - app_generator, - "agent_has_workflow_callable_active_snapshot", - is_publish_visible, - ) + def guard_result(**_kwargs): + events.append("queried") + return True, "query", annotation_reply - def test_success_chains_to_resolve_by_id(self, sqlite_session: Session): - bound_agent = _resolve_agent_agent() - snapshot = _resolve_agent_snapshot() - _persist(sqlite_session, bound_agent, snapshot) - app_model = _resolve_agent_app() + generator._run_input_guards.side_effect = guard_result - agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - ) + def publish(event, *_args): + if isinstance(event, QueueAnnotationReplyEvent): + events.append("publish") - assert agent is bound_agent - assert config_id == snapshot.id - assert config_version_kind == "snapshot" - assert soul.model is not None - - def test_unpublished_draft_still_resolves_active_snapshot( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - bound_agent = _resolve_agent_agent(published=False) - snapshot = _resolve_agent_snapshot() - _persist(sqlite_session, bound_agent, snapshot) - app_model = _resolve_agent_app() - monkeypatch.setattr(app_generator, "agent_has_workflow_callable_active_snapshot", lambda **_kwargs: True) - - agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - ) + queue_manager.publish.side_effect = publish - assert agent is bound_agent - assert config_id == snapshot.id - assert config_version_kind == "snapshot" - assert soul.prompt.system_prompt == "You are Iris." - - def test_existing_conversation_resolves_binding_snapshot_instead_of_latest_active_snapshot( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - bound_agent = _resolve_agent_agent(active_snapshot_id="snap-2") - pinned_snapshot = _resolve_agent_snapshot() - _persist(sqlite_session, bound_agent, pinned_snapshot) - conversation = _resolve_agent_conversation(binding_id="binding-1") - binding = _binding() - get_active_binding = MagicMock(return_value=binding) - validate_generation = MagicMock() - monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) - monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) - app_model = _resolve_agent_app() - - _, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - conversation=conversation, - ) + self._call(generator, mocker, queue_manager) - assert config_id == "snap-1" - assert config_version_kind == "snapshot" - assert soul.prompt.system_prompt == "You are Iris." - assert get_active_binding.call_args.kwargs["binding_id"] == "binding-1" - validate_generation.assert_called_once_with( - binding, - base_home_snapshot_id="home-1", - agent_config_version_id="snap-1", - agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, - ) + assert events == ["queried", "publish"] + runner.run.assert_not_called() - def test_existing_conversation_rejects_unavailable_binding( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - bound_agent = _resolve_agent_agent(active_snapshot_id="snap-active") - _persist(sqlite_session, bound_agent) - conversation = _resolve_agent_conversation(binding_id="binding-missing") - monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None)) - - with pytest.raises(AgentAppGeneratorError, match="Conversation participant Binding is unavailable"): - AgentAppGenerator()._resolve_agent( - _resolve_agent_app(), - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - conversation=conversation, - ) - - @pytest.mark.parametrize( - ("binding_home_id", "binding_version_kind", "snapshot_home_id"), - [ - ("home-binding", AgentConfigVersionKind.SNAPSHOT, "home-other"), - ("home-pinned", AgentConfigVersionKind.DRAFT, "home-pinned"), - ], + def test_generate_task_stopped_is_swallowed(self, generator, mocker: MockerFixture): + self._wire(generator, mocker, run_side_effect=GenerateTaskStoppedError()) + # Local stop intent is not a backend terminal; tracing finalization stays + # with the runner event that proves the terminal outcome. + mocker.patch( + "core.ops.base_trace_instance.application_tracing_hooks", + side_effect=AssertionError("stop intent must not invent a tracing terminal"), ) - def test_existing_conversation_generation_mismatch_does_not_fallback_to_active_snapshot( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - binding_home_id: str, - binding_version_kind: AgentConfigVersionKind, - snapshot_home_id: str, - ): - bound_agent = _resolve_agent_agent(active_snapshot_id="snap-active") - conversation = _resolve_agent_conversation(binding_id="binding-1") - binding = _binding( - home_snapshot_id=binding_home_id, - version_id="snap-pinned", - version_kind=binding_version_kind, - ) - pinned_snapshot = _resolve_agent_snapshot(snapshot_id="snap-pinned", home_snapshot_id=snapshot_home_id) - _persist(sqlite_session, bound_agent, pinned_snapshot) - monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) - - with pytest.raises(AgentWorkspaceBindingGenerationMismatchError): - AgentAppGenerator()._resolve_agent( - _resolve_agent_app(), - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - conversation=conversation, - ) - - def test_unpublished_imported_agent_is_not_available_to_public_runtime(self, sqlite_session: Session): - bound_agent = _resolve_agent_agent(source=AgentSource.IMPORTED, published=False) - _persist(sqlite_session, bound_agent) - app_model = _resolve_agent_app() - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - ) - - def test_never_published_agent_app_is_not_available_to_public_runtime(self, sqlite_session: Session): - bound_agent = _resolve_agent_agent(published=False) - _persist(sqlite_session, bound_agent) - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - AgentAppGenerator()._resolve_agent( - _resolve_agent_app(), - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - ) - - def test_unpublished_imported_agent_remains_available_to_debugger(self, sqlite_session: Session): - bound_agent = _resolve_agent_agent(source=AgentSource.IMPORTED, published=False) - draft = AgentConfigDraft( - id="draft-1", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DRAFT, - account_id=None, - draft_owner_key="", - base_snapshot_id="snap-1", - home_snapshot_id="home-1", - config_snapshot=AgentSoulConfig.model_validate(_SOUL_DICT), - ) - _persist(sqlite_session, bound_agent, draft) - app_model = _resolve_agent_app() - - agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.DEBUGGER, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - ) - - assert agent is bound_agent - assert config_id == draft.id - assert config_version_kind == "draft" - assert soul.prompt.system_prompt == "You are Iris." - - def test_agent_without_active_snapshot_raises_before_model_resolution(self, sqlite_session: Session): - bound_agent = _resolve_agent_agent(active_snapshot_id=None, published=False) - _persist(sqlite_session, bound_agent) - app_model = _resolve_agent_app() - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - ) - - def test_unbound_app_raises(self, sqlite_session: Session): - app_model = _resolve_agent_app() - with pytest.raises(AgentAppGeneratorError, match="has no bound Agent"): - AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_resolve_agent_account(), - session=sqlite_session, - ) + queue_manager = mocker.MagicMock() + self._call(generator, mocker, queue_manager) + queue_manager.publish_error.assert_not_called() + + def test_unexpected_error_is_published(self, generator: AgentAppGenerator, mocker: MockerFixture): + self._wire(generator, mocker, run_side_effect=ValueError("boom")) + queue_manager = mocker.MagicMock() + self._call(generator, mocker, queue_manager) + assert queue_manager.publish_error.called diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py index ce33a95dda5895..86ae6528e42848 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py @@ -9,7 +9,6 @@ from datetime import UTC, datetime from decimal import Decimal from typing import Any, override -from unittest.mock import MagicMock import pytest from dify_agent.protocol import ( @@ -49,9 +48,6 @@ FakeAgentBackendRunClient, FakeAgentBackendScenario, ) -from core.app.apps.agent_app import app_runner as app_runner_module -from core.app.apps.agent_app.app_runner import AgentAppRunner -from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext, AgentAppRuntimeRequestBuilder from core.app.apps.agent_app.session_store import AgentAppSessionScope, StoredAgentAppSession from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom @@ -62,20 +58,37 @@ QueueMessageEndEvent, QueueRetrieverResourcesEvent, ) -from core.workflow.nodes.agent_v2.dify_tools_builder import WorkflowAgentToolLayers from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.errors.invoke import InvokeRateLimitError from models.agent_config_entities import AgentSoulConfig from models.enums import ConversationFromSource from models.model import AppMode, Message, MessageAgentThought +from services.app.generation.adapters import agent_runner as app_runner_module +from services.app.generation.adapters.agent_request_builder import ( + AgentAppRuntimeBuildContext, + AgentAppRuntimeRequestBuilder, +) +from services.app.generation.adapters.agent_runner import AgentAppRunner +from services.app.generation.ports import MessageIdentity +from services.workflow.execution.adapters.agent_v2.dify_tools_builder import WorkflowAgentToolLayers +from tests.helpers.unit.model_factories import make_app, make_conversation @pytest.fixture(autouse=True) def bind_agent_dependencies(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: """Bind local runner dependencies without reaching external services.""" - monkeypatch.setattr(app_runner_module.db, "session", sqlite_session) + sqlite_session.add_all( + [ + make_app(), + make_conversation( + conversation_id="conv-1", app_id="app-1", inputs={}, from_source=ConversationFromSource.API + ), + _message_record(), + ] + ) + sqlite_session.commit() monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_context_window", + "services.app.generation.adapters.agent_request_builder.resolve_model_context_window", lambda **_kwargs: None, ) @@ -194,27 +207,6 @@ def stream_events( ) -class _UsagePausedClient(FakeAgentBackendRunClient): - def __init__(self) -> None: - super().__init__(scenario=FakeAgentBackendScenario.PAUSED) - - @override - def _events(self, run_id: str) -> tuple[RunEvent, ...]: - events = super()._events(run_id) - terminal = events[-1] - assert isinstance(terminal, RunSucceededEvent) - return ( - *events[:-1], - terminal.model_copy( - update={ - "data": terminal.data.model_copy( - update={"usage": AgentRunUsage(prompt_tokens=5, completion_tokens=3)} - ) - } - ), - ) - - class _RunLimitBindingLostFakeAgentBackendRunClient(FakeAgentBackendRunClient): @override def stream_events( @@ -651,8 +643,10 @@ def _runner( store: _FakeSessionStore, *, text_delta_debounce_seconds: float = 0, + app_records, ) -> AgentAppRunner: return AgentAppRunner( + records=app_records, request_builder=AgentAppRuntimeRequestBuilder( dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] ), @@ -724,12 +718,12 @@ def _saved_user_query(qm: _FakeQueueManager) -> str: return content -def test_successful_turn_publishes_chunk_and_message_end_and_saves_session() -> None: +def test_successful_turn_publishes_chunk_and_message_end_and_saves_session(*, app_records) -> None: client = FakeAgentBackendRunClient() # SUCCESS: output {"text": "hello agent"} store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) assert client.request is not None # One LLM chunk + one message-end, carrying the backend's answer text. @@ -750,11 +744,11 @@ def test_successful_turn_publishes_chunk_and_message_end_and_saves_session() -> assert saved_snapshot is not None -def test_turn_uses_resolved_backend_binding_before_backend_invocation() -> None: +def test_turn_uses_resolved_backend_binding_before_backend_invocation(*, app_records) -> None: client = FakeAgentBackendRunClient() store = _FakeSessionStore(binding_id="binding-2", backend_binding_ref="backend-binding-2") - _run(_runner(client, store), _FakeQueueManager()) + _run(_runner(client, store, app_records=app_records), _FakeQueueManager()) assert client.request is not None layers = {layer["name"]: layer for layer in client.request.model_dump(mode="json")["composition"]["layers"]} @@ -764,13 +758,13 @@ def test_turn_uses_resolved_backend_binding_before_backend_invocation() -> None: def test_successful_turn_routes_stream_text_to_agent_message_and_uses_terminal_output( - sqlite_session: Session, + sqlite_session: Session, *, app_records ) -> None: client = _StreamingFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -792,12 +786,12 @@ def test_successful_turn_routes_stream_text_to_agent_message_and_uses_terminal_o assert store.saved -def test_successful_turn_persists_usage_without_a_queue_consumer(sqlite_session: Session) -> None: - sqlite_session.add(_message_record()) - sqlite_session.flush() +def test_successful_turn_persists_usage_without_a_queue_consumer(sqlite_session: Session, *, app_records) -> None: + sqlite_session.merge(_message_record()) + sqlite_session.commit() _run( - _runner(_StreamingFakeAgentBackendRunClient(), _FakeSessionStore()), + _runner(_StreamingFakeAgentBackendRunClient(), _FakeSessionStore(), app_records=app_records), _FakeQueueManager(), ) @@ -811,12 +805,12 @@ def test_successful_turn_persists_usage_without_a_queue_consumer(sqlite_session: assert json.loads(message.message_metadata)["usage"]["total_tokens"] == 8 -def test_successful_turn_routes_single_agent_message_delta(sqlite_session: Session) -> None: +def test_successful_turn_routes_single_agent_message_delta(sqlite_session: Session, *, app_records) -> None: client = _StreamingSingleAgentMessageDeltaFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -829,7 +823,7 @@ def test_successful_turn_routes_single_agent_message_delta(sqlite_session: Sessi assert rows == [] -def test_thought_commit_failure_rolls_back_and_turn_continues(sqlite_session: Session) -> None: +def test_thought_commit_failure_rolls_back_and_turn_continues(sqlite_session: Session, *, app_records) -> None: rollback_events: list[Session] = [] should_fail = True @@ -842,29 +836,30 @@ def fail_first_commit(_session: Session) -> None: def record_rollback(session: Session) -> None: rollback_events.append(session) - event.listen(sqlite_session, "before_commit", fail_first_commit) - event.listen(sqlite_session, "after_rollback", record_rollback) + event.listen(Session, "before_commit", fail_first_commit) + event.listen(Session, "after_rollback", record_rollback) try: client = _StreamingSingleAgentMessageDeltaFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) finally: - event.remove(sqlite_session, "before_commit", fail_first_commit) - event.remove(sqlite_session, "after_rollback", record_rollback) + event.remove(Session, "before_commit", fail_first_commit) + event.remove(Session, "after_rollback", record_rollback) - assert rollback_events == [sqlite_session] + assert len(rollback_events) == 1 + assert rollback_events[0] is not sqlite_session assert _thought_rows(sqlite_session) == [] assert _llm_result(qm).message.content == "hello agent" -def test_successful_turn_with_null_terminal_output_publishes_empty_answer_not_literal_null() -> None: +def test_successful_turn_with_null_terminal_output_publishes_empty_answer_not_literal_null(*, app_records) -> None: client = _NullOutputFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -876,13 +871,13 @@ def test_successful_turn_with_null_terminal_output_publishes_empty_answer_not_li def test_successful_turn_with_stream_text_and_null_terminal_output_keeps_empty_message( - sqlite_session: Session, + sqlite_session: Session, *, app_records ) -> None: client = _StreamingTextNullOutputFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -896,12 +891,12 @@ def test_successful_turn_with_stream_text_and_null_terminal_output_keeps_empty_m assert rows[0].answer == "streamed answer" -def test_successful_turn_routes_agent_answer_to_agent_message(sqlite_session: Session) -> None: +def test_successful_turn_routes_agent_answer_to_agent_message(sqlite_session: Session, *, app_records) -> None: client = _AgentAnswerStreamingFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -921,14 +916,14 @@ def test_successful_turn_routes_agent_answer_to_agent_message(sqlite_session: Se def test_agent_message_deltas_are_debounced_to_agent_message( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, *, app_records ) -> None: monkeypatch.setattr(app_runner_module.time, "monotonic", _MonotonicClock(0.0, 0.2)) client = _StreamingFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store, text_delta_debounce_seconds=0.5), qm) + _run(_runner(client, store, app_records=app_records, text_delta_debounce_seconds=0.5), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -938,14 +933,12 @@ def test_agent_message_deltas_are_debounced_to_agent_message( assert rows == [] -def test_successful_turn_persists_thinking_and_tool_process_events( - sqlite_session: Session, -) -> None: +def test_successful_turn_persists_thinking_and_tool_process_events(sqlite_session: Session, *, app_records) -> None: client = _ProcessStreamingFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -963,15 +956,13 @@ def test_successful_turn_persists_thinking_and_tool_process_events( assert len(rows) == 2 -def test_streaming_turn_cancels_after_persisting_seen_agent_answer( - sqlite_session: Session, -) -> None: +def test_streaming_turn_cancels_after_persisting_seen_agent_answer(sqlite_session: Session, *, app_records) -> None: store = _FakeSessionStore() qm = _FakeQueueManager() client = _StreamingStopAfterFirstDeltaFakeAgentBackendRunClient(queue_manager=qm) with pytest.raises(GenerateTaskStoppedError): - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] @@ -984,11 +975,11 @@ def test_streaming_turn_cancels_after_persisting_seen_agent_answer( assert client.cancel_after == ["3-0"] -def test_tool_result_without_identity_does_not_attach_to_previous_tool( - sqlite_session: Session, -) -> None: +def test_tool_result_without_identity_does_not_attach_to_previous_tool(sqlite_session: Session, *, app_records) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1028,9 +1019,11 @@ def test_tool_result_without_identity_does_not_attach_to_previous_tool( assert rows[1].observation == "Knowledge base search results: browser skill" -def test_trusted_knowledge_receipts_reach_chat_sources_once(sqlite_session: Session) -> None: +def test_trusted_knowledge_receipts_reach_chat_sources_once(sqlite_session: Session, *, app_records) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1063,9 +1056,11 @@ def test_trusted_knowledge_receipts_reach_chat_sources_once(sqlite_session: Sess assert _thought_rows(sqlite_session)[0].tool == "shell_run" -def test_answer_suffix_trim_keeps_non_terminal_prefix(sqlite_session: Session) -> None: +def test_answer_suffix_trim_keeps_non_terminal_prefix(sqlite_session: Session, *, app_records) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1079,9 +1074,11 @@ def test_answer_suffix_trim_keeps_non_terminal_prefix(sqlite_session: Session) - assert rows[0].answer == "intermediate " -def test_tool_call_part_binds_late_call_id_to_delta_row(sqlite_session: Session) -> None: +def test_tool_call_part_binds_late_call_id_to_delta_row(sqlite_session: Session, *, app_records) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1138,9 +1135,11 @@ def test_tool_call_part_binds_late_call_id_to_delta_row(sqlite_session: Session) assert rows[0].observation == "Knowledge base search results: browser skill" -def test_thinking_after_tool_starts_new_snapshot_row(sqlite_session: Session) -> None: +def test_thinking_after_tool_starts_new_snapshot_row(sqlite_session: Session, *, app_records) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1194,11 +1193,11 @@ def test_thinking_after_tool_starts_new_snapshot_row(sqlite_session: Session) -> assert rows[1].tool_input == '{"cmd": "date"}' -def test_tool_result_without_call_id_matches_unique_open_tool_name( - sqlite_session: Session, -) -> None: +def test_tool_result_without_call_id_matches_unique_open_tool_name(sqlite_session: Session, *, app_records) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1239,10 +1238,12 @@ def test_tool_result_without_call_id_matches_unique_open_tool_name( def test_repeated_tool_calls_without_call_id_or_index_create_distinct_rows( - sqlite_session: Session, + sqlite_session: Session, *, app_records ) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1312,10 +1313,12 @@ def test_repeated_tool_calls_without_call_id_or_index_create_distinct_rows( def test_repeated_tool_calls_with_placeholder_call_id_and_reused_index_create_distinct_rows( - sqlite_session: Session, + sqlite_session: Session, *, app_records ) -> None: qm = _FakeQueueManager() recorder = app_runner_module._AgentProcessRecorder( + records=app_records, + conversation_id="conv-1", dify_context=_dify_ctx(), message_id="msg-1", queue_manager=qm, # type: ignore[arg-type] @@ -1362,38 +1365,38 @@ def test_repeated_tool_calls_with_placeholder_call_id_and_reused_index_create_di assert rows[1].observation == "out output" -def test_prior_session_snapshot_is_threaded_into_request() -> None: +def test_prior_session_snapshot_is_threaded_into_request(*, app_records) -> None: prior = _compatible_session_snapshot() client = FakeAgentBackendRunClient() store = _FakeSessionStore(loaded=prior) qm = _FakeQueueManager() - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) assert client.request is not None assert client.request.session_snapshot is prior -def test_snapshot_with_removed_optional_prompt_is_forwarded_to_runtime() -> None: +def test_snapshot_with_removed_optional_prompt_is_forwarded_to_runtime(*, app_records) -> None: compatible = _compatible_session_snapshot() snapshot = SessionSnapshot( layers={name: state for name, state in compatible.layers.items() if name != "agent_soul_prompt"} ) client = FakeAgentBackendRunClient() store = _FakeSessionStore(loaded=snapshot) - _run(_runner(client, store), _FakeQueueManager()) + _run(_runner(client, store, app_records=app_records), _FakeQueueManager()) assert client.request is not None assert client.request.session_snapshot is snapshot assert len(store.saved) == 1 -def test_debug_session_scope_can_reuse_conversation_across_config_snapshots() -> None: +def test_debug_session_scope_can_reuse_conversation_across_config_snapshots(*, app_records) -> None: prior = _compatible_session_snapshot() client = FakeAgentBackendRunClient() store = _FakeSessionStore(loaded=prior) qm = _FakeQueueManager() - _runner(client, store).run( + _runner(client, store, app_records=app_records).run( dify_context=_dify_ctx(), agent_id="agent-1", agent_config_snapshot_id="snap-new", @@ -1413,24 +1416,27 @@ def test_debug_session_scope_can_reuse_conversation_across_config_snapshots() -> assert store.saved[0][0].agent_config_snapshot_id == "snap-new" -def test_failed_run_raises_agent_backend_error() -> None: +def test_failed_run_raises_agent_backend_error(*, app_records) -> None: client = FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario.FAILED) store = _FakeSessionStore() qm = _FakeQueueManager() with pytest.raises(AgentBackendRunFailedError, match="fake failure .*agent_run_id=fake-run-1"): - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) # No message-end on failure; post-exit session state is still saved. assert not [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] assert store.saved[0][2] == SessionSnapshot(layers={}) -def test_failed_run_persists_partial_usage(sqlite_session: Session) -> None: - sqlite_session.add(_message_record()) - sqlite_session.flush() +def test_failed_run_persists_partial_usage(sqlite_session: Session, *, app_records) -> None: + sqlite_session.merge(_message_record()) + sqlite_session.commit() with pytest.raises(AgentBackendRunFailedError, match="failed after model calls"): - _run(_runner(_UsageFailedClient(), _FakeSessionStore()), _FakeQueueManager()) + _run( + _runner(_UsageFailedClient(), _FakeSessionStore(), app_records=app_records), + _FakeQueueManager(), + ) sqlite_session.expire_all() message = sqlite_session.get(Message, "msg-1") @@ -1443,22 +1449,24 @@ def test_failed_run_persists_partial_usage(sqlite_session: Session) -> None: assert json.loads(message.message_metadata)["usage"]["total_tokens"] == 12 -def test_partial_usage_persistence_ignores_missing_message() -> None: - AgentAppRunner._persist_message_usage( - message_id="missing", +def test_partial_usage_persistence_ignores_missing_message(app_records) -> None: + _runner(FakeAgentBackendRunClient(), _FakeSessionStore(), app_records=app_records)._persist_message_usage( + identity=MessageIdentity("tenant-1", "app-1", "conv-1", "missing"), usage=LLMUsage.from_metadata({"prompt_tokens": 2, "completion_tokens": 1}), ) @pytest.mark.parametrize("metadata", ["{", "[]"]) -def test_partial_usage_persistence_recovers_invalid_metadata(metadata: str, sqlite_session: Session) -> None: +def test_partial_usage_persistence_recovers_invalid_metadata( + metadata: str, sqlite_session: Session, app_records +) -> None: message = _message_record() message.message_metadata = metadata - sqlite_session.add(message) - sqlite_session.flush() + sqlite_session.merge(message) + sqlite_session.commit() - AgentAppRunner._persist_message_usage( - message_id=message.id, + _runner(FakeAgentBackendRunClient(), _FakeSessionStore(), app_records=app_records)._persist_message_usage( + identity=MessageIdentity("tenant-1", "app-1", "conv-1", message.id), usage=LLMUsage.from_metadata({"prompt_tokens": 2, "completion_tokens": 1}), ) @@ -1469,21 +1477,27 @@ def test_partial_usage_persistence_recovers_invalid_metadata(metadata: str, sqli assert json.loads(persisted.message_metadata)["usage"]["total_tokens"] == 3 -def test_partial_usage_persistence_rolls_back_database_error(monkeypatch: pytest.MonkeyPatch) -> None: - session = MagicMock() - session.get.side_effect = RuntimeError("database unavailable") - monkeypatch.setattr(app_runner_module.db, "session", session) +def test_partial_usage_persistence_rolls_back_database_error(app_records, sqlite_session: Session) -> None: + engine = sqlite_session.get_bind() - AgentAppRunner._persist_message_usage( - message_id="msg-1", - usage=LLMUsage.from_metadata({"prompt_tokens": 2, "completion_tokens": 1}), - ) + def fail(_conn, _cursor, statement, _params, _ctx, _many): + if statement.startswith("UPDATE messages"): + raise RuntimeError("database unavailable") - session.rollback.assert_called_once() + event.listen(engine, "before_cursor_execute", fail) + try: + _runner(FakeAgentBackendRunClient(), _FakeSessionStore(), app_records=app_records)._persist_message_usage( + identity=MessageIdentity("tenant-1", "app-1", "conv-1", "msg-1"), + usage=LLMUsage.from_metadata({"prompt_tokens": 2, "completion_tokens": 1}), + ) + finally: + event.remove(engine, "before_cursor_execute", fail) + sqlite_session.expire_all() + assert sqlite_session.get(Message, "msg-1").message_tokens == 0 @pytest.mark.parametrize("outcome", ["failed", "stopped"]) -def test_snapshot_save_failure_preserves_original_app_outcome(outcome: str) -> None: +def test_snapshot_save_failure_preserves_original_app_outcome(outcome: str, *, app_records) -> None: store = _ExplodingSessionStore() queue_manager: _FakeQueueManager = _FakeQueueManager() if outcome == "failed" else _StoppedQueueManager() client = FakeAgentBackendRunClient( @@ -1492,7 +1506,7 @@ def test_snapshot_save_failure_preserves_original_app_outcome(outcome: str) -> N expected_error = AgentBackendRunFailedError if outcome == "failed" else GenerateTaskStoppedError with pytest.raises(expected_error, match="fake failure" if outcome == "failed" else None): - _run(_runner(client, store), queue_manager) + _run(_runner(client, store, app_records=app_records), queue_manager) assert store.save_attempts == [SessionSnapshot(layers={})] @@ -1502,24 +1516,23 @@ def test_snapshot_save_failure_preserves_original_app_outcome(outcome: str) -> N [("failed", AgentBackendRunFailedError), ("cancelled", AgentBackendError)], ) def test_terminal_without_snapshot_preserves_prior_app_session_without_write( - terminal_type: str, - expected_error: type[Exception], + terminal_type: str, expected_error: type[Exception], *, app_records ) -> None: store = _FakeSessionStore() client = _TerminalWithoutSnapshotFakeAgentBackendRunClient(terminal_type=terminal_type) with pytest.raises(expected_error): - _run(_runner(client, store), _FakeQueueManager()) + _run(_runner(client, store, app_records=app_records), _FakeQueueManager()) assert store.saved == [] -def test_failed_run_prefers_run_failure_type_over_binding_lost_reason() -> None: +def test_failed_run_prefers_run_failure_type_over_binding_lost_reason(*, app_records) -> None: client = _RunLimitBindingLostFakeAgentBackendRunClient() store = _FakeSessionStore() with pytest.raises(AgentBackendRunFailedError) as raised: - _run(_runner(client, store), _FakeQueueManager()) + _run(_runner(client, store, app_records=app_records), _FakeQueueManager()) assert raised.value.error_type is RunFailureType.AGENT_RUN_LIMIT_EXCEEDED assert raised.value.reason == "binding_lost" @@ -1580,26 +1593,29 @@ def test_agent_backend_failure_to_exception_prefers_run_failure_type_over_known_ } -def test_stopped_task_waits_for_cancelled_snapshot_and_saves_session() -> None: +def test_stopped_task_waits_for_cancelled_snapshot_and_saves_session(*, app_records) -> None: client = _RecordingFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _StoppedQueueManager() with pytest.raises(GenerateTaskStoppedError): - _run(_runner(client, store), qm) + _run(_runner(client, store, app_records=app_records), qm) assert client.cancelled_run_ids == ["fake-run-1"] assert len(store.saved) == 1 assert store.saved[0][2] == SessionSnapshot(layers={}) -def test_stopped_task_persists_partial_usage(sqlite_session: Session) -> None: - sqlite_session.add(_message_record()) - sqlite_session.flush() +def test_stopped_task_persists_partial_usage(sqlite_session: Session, *, app_records) -> None: + sqlite_session.merge(_message_record()) + sqlite_session.commit() client = _UsageCancellationClient() with pytest.raises(GenerateTaskStoppedError): - _run(_runner(client, _FakeSessionStore()), _StoppedQueueManager()) + _run( + _runner(client, _FakeSessionStore(), app_records=app_records), + _StoppedQueueManager(), + ) sqlite_session.expire_all() message = sqlite_session.get(Message, "msg-1") @@ -1611,12 +1627,12 @@ def test_stopped_task_persists_partial_usage(sqlite_session: Session) -> None: assert json.loads(message.message_metadata)["usage"]["total_tokens"] == 21 -def test_cancel_and_wait_failure_preserves_stopped_app_outcome() -> None: +def test_cancel_and_wait_failure_preserves_stopped_app_outcome(*, app_records) -> None: client = _CancelAndWaitFailingClient() store = _FakeSessionStore() with pytest.raises(GenerateTaskStoppedError): - _run(_runner(client, store), _StoppedQueueManager()) + _run(_runner(client, store, app_records=app_records), _StoppedQueueManager()) assert client.cancel_after == [None] assert store.saved == [] diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py b/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py new file mode 100644 index 00000000000000..eedb4fd4b50e9c --- /dev/null +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py @@ -0,0 +1,290 @@ +"""Unit tests for AgentAppGenerator._run_input_guards. + +The guards apply content moderation + annotation reply before the Agent backend +call, mirroring the EasyUI chat runners. They can short-circuit the turn (a +blocked/preset moderation answer or a matched annotation) and publish a direct +answer, or pass through a possibly moderation-sanitized query. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import override +from unittest.mock import create_autospec + +import pytest + +import core.moderation.input_moderation as input_moderation_mod +from core.app.app_config.entities import ( + AppAdditionalFeatures, + EasyUIBasedAppConfig, + EasyUIBasedAppModelConfigFrom, + ModelConfigEntity, + PromptTemplateEntity, +) +from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom +from core.app.entities.app_invoke_entities import AgentAppGenerateEntity, InvokeFrom +from core.app.entities.queue_entities import ( + AppQueueEvent, + QueueLLMChunkEvent, + QueueMessageEndEvent, +) +from core.moderation.base import ModerationError +from core.moderation.input_moderation import InputModeration +from models.annotation_reply import AnnotationReplies, AnnotationReply +from models.model import App, AppMode, Message +from models.tool_runtime_contracts import WorkflowToolQueries +from repositories.app.generation_repository import AppGenerationRepository +from repositories.tools.provider_repository import ToolProviderRepository +from services.app.generation.adapters.agent_app import AgentAppGenerator +from services.app.generation.agent_config import AgentAppConfigurations +from tests.helpers.unit.core.model_fixtures import make_model_config +from tests.helpers.unit.model_factories import make_app, make_message + + +class _FakeQueueManager(AppQueueManager): + def __init__(self) -> None: + self.events: list[AppQueueEvent] = [] + + @override + def publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: + del pub_from + self.events.append(event) + + @override + def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: + del pub_from + self.events.append(event) + + +def _app_config() -> EasyUIBasedAppConfig: + return EasyUIBasedAppConfig( + tenant_id="tenant-1", + app_id="app-1", + app_mode=AppMode.AGENT, + app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, + app_model_config_id=None, + app_model_config_dict={}, + model=ModelConfigEntity(provider="test-provider", model="gpt-4o-mini", mode="chat"), + prompt_template=PromptTemplateEntity( + prompt_type=PromptTemplateEntity.PromptType.SIMPLE, + simple_prompt_template="Hello", + ), + additional_features=AppAdditionalFeatures(), + variables=[], + ) + + +def _make_entity(query: str = "hello") -> AgentAppGenerateEntity: + return AgentAppGenerateEntity( + task_id="task-1", + app_config=_app_config(), + model_conf=make_model_config(provider="test-provider", model="gpt-4o-mini", mode="chat"), + file_upload_config=None, + conversation_id=None, + inputs={}, + query=query, + files=[], + parent_message_id=None, + trace_manager=None, + user_id="user-1", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + call_depth=0, + agent_id="agent-1", + agent_config_snapshot_id="snapshot-1", + agent_config_version_kind="snapshot", + agent_session_scope_config_version_id="snapshot-1", + agent_llm_gateway_enabled=True, + ) + + +def _app() -> App: + return make_app( + app_id="app-1", + tenant_id="tenant-1", + name="Agent App", + mode=AppMode.AGENT, + icon_type=None, + ) + + +def _message() -> Message: + return make_message(message_id="msg-1") + + +ModerationResult = tuple[bool, Mapping[str, object], str] + + +def _patch_moderation( + monkeypatch: pytest.MonkeyPatch, + *, + returns: ModerationResult | None = None, + raises: Exception | None = None, +) -> None: + def check(_moderation: InputModeration, **_kwargs: object) -> ModerationResult: + if raises is not None: + raise raises + if returns is None: + raise AssertionError("a moderation result is required when no exception is configured") + return returns + + monkeypatch.setattr(input_moderation_mod.InputModeration, "check", check) + + +def _patch_annotation( + monkeypatch: pytest.MonkeyPatch, + annotations: AnnotationReplies, + *, + reply: AnnotationReply | None, +) -> None: + def query(**_kwargs: object) -> AnnotationReply | None: + return reply + + monkeypatch.setattr(annotations, "query", query) + + +def _answer_text(events: list[AppQueueEvent]) -> str: + end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) + assert end.llm_result is not None + content = end.llm_result.message.content + assert isinstance(content, str) + return content + + +def _saved_user_query(events: list[AppQueueEvent]) -> str: + end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) + assert end.llm_result is not None + prompt_messages = end.llm_result.prompt_messages + assert len(prompt_messages) == 1 + content = prompt_messages[0].content + assert isinstance(content, str) + return content + + +class TestRunInputGuards: + def test_no_guards_passes_through( + self, + tool_providers: ToolProviderRepository, + monkeypatch: pytest.MonkeyPatch, + *, + annotation_replies: AnnotationReplies, + app_records: AppGenerationRepository, + ) -> None: + _patch_moderation(monkeypatch, returns=(False, {}, "hello")) + _patch_annotation(monkeypatch, annotation_replies, reply=None) + qm = _FakeQueueManager() + + handled, query, annotation_reply = AgentAppGenerator( + agent_configs=create_autospec(AgentAppConfigurations, instance=True, spec_set=True), + annotations=annotation_replies, + records=app_records, + tool_providers=tool_providers, + workflow_queries=create_autospec(WorkflowToolQueries, instance=True, spec_set=True), + )._run_input_guards( + application_generate_entity=_make_entity("hello"), + app_model=_app(), + message=_message(), + queue_manager=qm, + ) + + assert handled is False + assert query == "hello" + assert annotation_reply is None + assert qm.events == [] + + def test_moderation_override_sanitizes_query( + self, + tool_providers: ToolProviderRepository, + monkeypatch: pytest.MonkeyPatch, + *, + annotation_replies: AnnotationReplies, + app_records: AppGenerationRepository, + ) -> None: + _patch_moderation(monkeypatch, returns=(True, {}, "[redacted]")) + _patch_annotation(monkeypatch, annotation_replies, reply=None) + qm = _FakeQueueManager() + + handled, query, annotation_reply = AgentAppGenerator( + agent_configs=create_autospec(AgentAppConfigurations, instance=True, spec_set=True), + annotations=annotation_replies, + records=app_records, + tool_providers=tool_providers, + workflow_queries=create_autospec(WorkflowToolQueries, instance=True, spec_set=True), + )._run_input_guards( + application_generate_entity=_make_entity("leak my secret"), + app_model=_app(), + message=_message(), + queue_manager=qm, + ) + + assert handled is False + assert query == "[redacted]" + assert annotation_reply is None + assert qm.events == [] + + def test_moderation_block_short_circuits( + self, + tool_providers: ToolProviderRepository, + monkeypatch: pytest.MonkeyPatch, + *, + annotation_replies: AnnotationReplies, + app_records: AppGenerationRepository, + ) -> None: + _patch_moderation(monkeypatch, raises=ModerationError("blocked preset answer")) + _patch_annotation(monkeypatch, annotation_replies, reply=None) + qm = _FakeQueueManager() + + handled, _, annotation_reply = AgentAppGenerator( + agent_configs=create_autospec(AgentAppConfigurations, instance=True, spec_set=True), + annotations=annotation_replies, + records=app_records, + tool_providers=tool_providers, + workflow_queries=create_autospec(WorkflowToolQueries, instance=True, spec_set=True), + )._run_input_guards( + application_generate_entity=_make_entity("forbidden"), + app_model=_app(), + message=_message(), + queue_manager=qm, + ) + + assert handled is True + assert annotation_reply is None + assert any(isinstance(e, QueueLLMChunkEvent) for e in qm.events) + assert _answer_text(qm.events) == "blocked preset answer" + assert _saved_user_query(qm.events) == "forbidden" + + def test_annotation_hit_short_circuits( + self, + tool_providers: ToolProviderRepository, + monkeypatch: pytest.MonkeyPatch, + *, + annotation_replies: AnnotationReplies, + app_records: AppGenerationRepository, + ) -> None: + _patch_moderation(monkeypatch, returns=(False, {}, "what is your name")) + annotation = AnnotationReply( + id="anno-1", + content="I am the annotated Iris.", + ) + _patch_annotation(monkeypatch, annotation_replies, reply=annotation) + qm = _FakeQueueManager() + + handled, _, annotation_reply = AgentAppGenerator( + agent_configs=create_autospec(AgentAppConfigurations, instance=True, spec_set=True), + annotations=annotation_replies, + records=app_records, + tool_providers=tool_providers, + workflow_queries=create_autospec(WorkflowToolQueries, instance=True, spec_set=True), + )._run_input_guards( + application_generate_entity=_make_entity("what is your name"), + app_model=_app(), + message=_message(), + queue_manager=qm, + ) + + assert handled is True + assert annotation_reply is not None + assert annotation_reply.id == "anno-1" + assert qm.events == [] diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py index a9b01d9aca1b53..bb52c305181e83 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py @@ -28,19 +28,18 @@ AgentBackendRunRequestBuilder, ) from clients.agent_backend.request_builder import DIFY_SHELL_LAYER_ID -from core.app.apps.agent_app.runtime_request_builder import ( - AgentAppRuntimeBuildContext, - AgentAppRuntimeRequestBuilder, - AgentAppRuntimeRequestBuildError, -) from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom from core.ops.base_trace_instance import bind_application_tracing_hooks from core.workflow.file_reference import build_file_reference -from core.workflow.nodes.agent_v2.dify_tools_builder import WorkflowAgentDifyToolsBuilder from graphon.file import File, FileTransferMethod, FileType from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent from graphon.model_runtime.entities.model_entities import ModelFeature from models.agent_config_entities import AgentSoulConfig +from services.app.generation.adapters.agent_request_builder import ( + AgentAppRuntimeBuildContext, + AgentAppRuntimeRequestBuilder, + AgentAppRuntimeRequestBuildError, +) from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema @@ -61,7 +60,7 @@ def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.fixture(autouse=True) def _no_runtime_agent_skills(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.load_runtime_agent_skill_configs", + "services.app.generation.adapters.agent_request_builder.load_runtime_agent_skill_configs", lambda **_kwargs: [], ) @@ -75,7 +74,7 @@ def resolve(*, run_context: object, provider_name: str, model_name: str) -> int: return 32_768 monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_context_window", + "services.app.generation.adapters.agent_request_builder.resolve_model_context_window", resolve, ) return calls @@ -160,7 +159,7 @@ def prepare(value: AgentAppRuntimeBuildContext) -> ApplicationTracingContext: bind_application_tracing_hooks(app, hooks) with app.app_context(): - result = AgentAppRuntimeRequestBuilder(dify_tools_builder=WorkflowAgentDifyToolsBuilder()).build(context) + result = AgentAppRuntimeRequestBuilder(dify_tools_builder=_NoToolsBuilder()).build(context) # type: ignore[arg-type] layer = next(item for item in result.request.composition.layers if item.name == "execution_context") assert DifyExecutionContextLayerConfig.model_validate(layer.config).tracing == tracing redacted = next( @@ -175,7 +174,7 @@ def __init__(self): self.knowledge_fs_spaces = () def build_layers(self, **kwargs): - self.knowledge_fs_spaces = kwargs["knowledge_fs_spaces"] + self.knowledge_fs_spaces = kwargs.get("knowledge_fs_spaces", ()) return SimpleNamespace(plugin_tools=None, core_tools=None, exposed_tool_names=lambda: []) @@ -343,7 +342,7 @@ def test_build_maps_soul_to_run_request(self, model_context_window_calls: list[t def test_build_sends_images_directly_to_vision_model(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_supports_vision", + "services.app.generation.adapters.agent_request_builder.resolve_model_supports_vision", lambda **_kwargs: True, ) prompt_content_calls: list[tuple[File, ImagePromptMessageContent.DETAIL | None]] = [] @@ -363,7 +362,7 @@ def to_prompt_message_content( ) monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.file_manager.to_prompt_message_content", + "services.app.generation.adapters.agent_request_builder.file_manager.to_prompt_message_content", to_prompt_message_content, ) builder = AgentAppRuntimeRequestBuilder(dify_tools_builder=_NoToolsBuilder()) # type: ignore[arg-type] @@ -390,7 +389,7 @@ def to_prompt_message_content( def test_build_keeps_image_locator_for_non_vision_model(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_supports_vision", + "services.app.generation.adapters.agent_request_builder.resolve_model_supports_vision", lambda **_kwargs: False, ) builder = AgentAppRuntimeRequestBuilder(dify_tools_builder=_NoToolsBuilder()) # type: ignore[arg-type] @@ -417,7 +416,7 @@ def test_knowledge_image_search_retains_named_upload_references( ) -> None: apply_config_overrides(monkeypatch, AGENT_SHELL_ENABLED=True) monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_supports_vision", + "services.app.generation.adapters.agent_request_builder.resolve_model_supports_vision", lambda **_kwargs: vision, ) model_factory = RecordingModelFactory(schema=llm_schema(features=[ModelFeature.VISION] if vision else [])) @@ -434,7 +433,8 @@ def image_content(file: File, **_kwargs: object) -> ImagePromptMessageContent: ) monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.file_manager.to_prompt_message_content", image_content + "services.app.generation.adapters.agent_request_builder.file_manager.to_prompt_message_content", + image_content, ) soul = _soul_with_model() soul.knowledge = AgentSoulConfig.model_validate( @@ -496,11 +496,11 @@ def image_content(file: File, **_kwargs: object) -> ImagePromptMessageContent: def test_build_preserves_inline_base64_transport_for_vision_model(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_supports_vision", + "services.app.generation.adapters.agent_request_builder.resolve_model_supports_vision", lambda **_kwargs: True, ) monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.file_manager.to_prompt_message_content", + "services.app.generation.adapters.agent_request_builder.file_manager.to_prompt_message_content", lambda *_args, **_kwargs: ImagePromptMessageContent( format="png", base64_data="aW1hZ2UtYnl0ZXM=", @@ -522,11 +522,11 @@ def test_build_preserves_inline_base64_transport_for_vision_model(self, monkeypa def test_build_keeps_non_image_locator_when_vision_image_is_direct(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_supports_vision", + "services.app.generation.adapters.agent_request_builder.resolve_model_supports_vision", lambda **_kwargs: True, ) monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.file_manager.to_prompt_message_content", + "services.app.generation.adapters.agent_request_builder.file_manager.to_prompt_message_content", lambda *_args, **_kwargs: ImagePromptMessageContent( format="png", url="https://files.example.com/earth.png", @@ -937,7 +937,7 @@ def test_config_layer_for_build_draft_marks_config_writable(self): def test_config_layer_includes_bound_workspace_skills(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.load_runtime_agent_skill_configs", + "services.app.generation.adapters.agent_request_builder.load_runtime_agent_skill_configs", lambda **_kwargs: [ DifyConfigSkillConfig( name="workspace-skill", diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py b/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py index 47c978263e9858..ed13fffb5a0f12 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py @@ -12,6 +12,7 @@ ) from models.enums import ConversationFromSource from models.model import App, AppMode, Conversation +from repositories.agent_workspace_repository import AgentWorkspaceRepository from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService @@ -137,7 +138,7 @@ def test_load_or_create_uses_exact_caller_binding( create = MagicMock() store = AgentAppWorkspaceStore() monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding) - monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", MagicMock()) + monkeypatch.setattr(AgentWorkspaceRepository, "validate_binding_generation", MagicMock()) monkeypatch.setattr(AgentWorkspaceService, "create_binding", create) stored = store.load_or_create(_scope()) diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py index e8d66310c26b4a..d7308d779451cf 100644 --- a/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py @@ -7,10 +7,8 @@ import pytest from flask import Flask from pydantic import ValidationError -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker -from core.app.apps.agent_chat.app_generator import AgentChatAppGenerator -from core.app.apps.agent_chat.app_runner import AgentChatAppRunner from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import ( @@ -20,10 +18,15 @@ from core.app.entities.queue_entities import QueueErrorEvent from core.app.file_access import get_current_file_access_scope from core.ops.ops_trace_manager import TraceQueueManager +from extensions.application_services.workflow import WorkflowExecutionDependencies from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from models import Account +from models.annotation_reply import AnnotationReplies +from models.enums import ConversationFromSource from models.model import App, AppMode, AppModelConfig, Conversation, Message -from services.workflow_run_agg import WorkflowRunAgg +from repositories.app.generation_repository import AppGenerationRepository +from services.app.generation.adapters.agent_chat import AgentChatAppGenerator +from services.app.generation.adapters.agent_chat_runner import AgentChatAppRunner from tests.helpers.unit.config_override import apply_config_overrides @@ -69,13 +72,24 @@ def generator( agent_queue_manager: MessageBasedAppQueueManager, agent_records: tuple[Conversation, Message], monkeypatch: pytest.MonkeyPatch, + app_records: AppGenerationRepository, + sqlite_session_factory: sessionmaker[Session], + *, + annotation_replies: AnnotationReplies, + workflow_runtime: WorkflowExecutionDependencies, ) -> AgentChatAppGenerator: assert agent_runtime_app.name assert agent_queue_manager is not None assert agent_records[0].id == "conv" # Keep trace configuration lookup and initialization real, but do not schedule background timers. monkeypatch.setattr(TraceQueueManager, "start_timer", lambda _manager: None) - return AgentChatAppGenerator(execution_driver=WorkflowRunAgg.run) + return AgentChatAppGenerator( + dataset_tools=workflow_runtime.dataset_tools, + annotations=annotation_replies, + records=app_records, + tool_sessions=sqlite_session_factory, + workflow_runtime=workflow_runtime, + ) @pytest.fixture @@ -187,6 +201,8 @@ def test_generate_success_with_debugger_override( user = _account() invoke_from = InvokeFrom.DEBUGGER + agent_records[0].from_source = ConversationFromSource.CONSOLE + sqlite_session.commit() get_model_config = CallRecorder(app_model_config) prepare_inputs = CallRecorder({"x": 1}) @@ -198,23 +214,20 @@ def test_generate_success_with_debugger_override( generator._handle_response = handle_response # type: ignore[method-assign] monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.AgentChatAppConfigManager.config_validate", + "services.app.generation.adapters.agent_chat.AgentChatAppConfigManager.config_validate", CallRecorder({"validated": True}), ) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.AgentChatAppConfigManager.get_app_config", + "services.app.generation.adapters.agent_chat.AgentChatAppConfigManager.get_app_config", CallRecorder(agent_generate_entity.app_config), ) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.ModelConfigConverter.convert", + "services.app.generation.adapters.agent_chat.ModelConfigConverter.convert", CallRecorder(agent_generate_entity.model_conf), ) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.FileUploadConfigManager.convert", CallRecorder(None) - ) - monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.ConversationService.get_conversation", - CallRecorder(agent_records[0]), + "services.app.generation.adapters.agent_chat.FileUploadConfigManager.convert", + CallRecorder(None), ) runner_calls: list[dict[str, object]] = [] @@ -223,7 +236,7 @@ def observe_runner(_runner: AgentChatAppRunner, **kwargs: object) -> None: monkeypatch.setattr(AgentChatAppRunner, "run", observe_runner) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.AgentChatAppGenerateResponseConverter.convert", + "services.app.generation.adapters.agent_chat.AgentChatAppGenerateResponseConverter.convert", CallRecorder({"result": "ok"}), ) @@ -247,8 +260,8 @@ def observe_runner(_runner: AgentChatAppRunner, **kwargs: object) -> None: ) assert result == {"result": "ok"} - assert get_model_config.calls[-1].kwargs["session"] is session - assert init_records.calls[-1].kwargs["session"] is session + assert "session" not in get_model_config.calls[-1].kwargs + assert "session" not in init_records.calls[-1].kwargs assert len(scheduled_threads) == 1 thread = scheduled_threads[0] assert type(thread) is threading.Thread @@ -295,21 +308,18 @@ def test_generate_without_file_config( monkeypatch.setattr(AppModelConfig, "to_dict", to_dict) load_annotation_reply_config = CallRecorder(annotation_reply) - monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.load_annotation_reply_config", - load_annotation_reply_config, - ) + monkeypatch.setattr(generator._records, "annotation_config", load_annotation_reply_config) get_app_config = CallRecorder(agent_generate_entity.app_config) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.AgentChatAppConfigManager.get_app_config", + "services.app.generation.adapters.agent_chat.AgentChatAppConfigManager.get_app_config", get_app_config, ) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.ModelConfigConverter.convert", + "services.app.generation.adapters.agent_chat.ModelConfigConverter.convert", CallRecorder(agent_generate_entity.model_conf), ) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.FileUploadConfigManager.convert", + "services.app.generation.adapters.agent_chat.FileUploadConfigManager.convert", CallRecorder(None), ) runner_calls: list[dict[str, object]] = [] @@ -320,7 +330,7 @@ def observe_runner(_runner: AgentChatAppRunner, **kwargs: object) -> None: monkeypatch.setattr(AgentChatAppRunner, "run", observe_runner) monkeypatch.setattr( - "core.app.apps.agent_chat.app_generator.AgentChatAppGenerateResponseConverter.convert", + "services.app.generation.adapters.agent_chat.AgentChatAppGenerateResponseConverter.convert", CallRecorder({"result": "ok"}), ) @@ -337,7 +347,9 @@ def observe_runner(_runner: AgentChatAppRunner, **kwargs: object) -> None: ) assert result == {"result": "ok"} - assert load_annotation_reply_config.calls == [RecordedCall(args=(session, "app1"), kwargs={})] + assert load_annotation_reply_config.calls == [ + RecordedCall(args=(), kwargs={"tenant_id": "tenant", "app_id": "app1"}) + ] assert to_dict.calls == [RecordedCall(args=(), kwargs={"annotation_reply": annotation_reply})] assert get_app_config.calls[-1].kwargs["annotation_reply"] is annotation_reply assert len(scheduled_threads) == 1 @@ -420,7 +432,7 @@ def test_generate_worker_logs_value_error_when_debug( _raise_in_runner(monkeypatch, ValueError("bad")) apply_config_overrides(monkeypatch, DEBUG=True) - with caplog.at_level(logging.ERROR, logger="core.app.apps.agent_chat.app_generator"): + with caplog.at_level(logging.ERROR, logger="services.app.generation.adapters.agent_chat"): generator._generate_worker( flask_app=agent_runtime_app, context=contextvars.copy_context(), diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_runner.py index a8ea37696c8175..aabcf9a3da0c97 100644 --- a/api/tests/unit_tests/core/app/apps/agent_chat/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_runner.py @@ -1,16 +1,12 @@ from collections.abc import Generator, Iterator, Mapping +from unittest.mock import create_autospec import pytest -from sqlalchemy import event, select +from sqlalchemy import select from sqlalchemy.orm import Session -from core.agent.cot_agent_runner import CotAgentRunner -from core.agent.cot_chat_agent_runner import CotChatAgentRunner -from core.agent.cot_completion_agent_runner import CotCompletionAgentRunner from core.agent.entities import AgentEntity -from core.agent.fc_agent_runner import FunctionCallAgentRunner from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig -from core.app.apps.agent_chat.app_runner import AgentChatAppRunner from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueLLMChunkEvent, QueueMessageEndEvent @@ -18,11 +14,19 @@ from core.moderation.input_moderation import InputModeration from core.plugin.entities.plugin_daemon import PluginLLMNumTokensResponse from core.plugin.impl.model import PluginModelClient +from extensions.application_services.workflow import WorkflowExecutionDependencies from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage from graphon.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey +from models.annotation_reply import AnnotationReplies from models.model import App, Conversation, Message, MessageAgentThought, MessageAnnotation -from services.workflow_run_agg import WorkflowRunAgg +from repositories.app.generation_repository import AppGenerationRepository +from services.agent.chat.cot_chat_runner import CotChatAgentRunner +from services.agent.chat.cot_completion_runner import CotCompletionAgentRunner +from services.agent.chat.cot_runner import CotAgentRunner +from services.agent.chat.function_call_runner import FunctionCallAgentRunner +from services.agent.chat.ports import AgentToolInvoker +from services.app.generation.adapters.agent_chat_runner import AgentChatAppRunner def _assert_direct_output(agent_queue_manager: MessageBasedAppQueueManager, text: str) -> None: @@ -36,36 +40,33 @@ def _assert_direct_output(agent_queue_manager: MessageBasedAppQueueManager, text @pytest.fixture -def runner(agent_records: tuple[Conversation, Message]) -> AgentChatAppRunner: +def runner( + agent_records: tuple[Conversation, Message], + *, + app_records: AppGenerationRepository, + workflow_runtime: WorkflowExecutionDependencies, +) -> AgentChatAppRunner: assert agent_records[0].id == "conv" - return AgentChatAppRunner(execution_driver=WorkflowRunAgg.run) + return AgentChatAppRunner( + dataset_tools=workflow_runtime.dataset_tools, + records=app_records, + tool_invoker=create_autospec(AgentToolInvoker, instance=True, spec_set=True), + ) -def _records(session: Session) -> tuple[Conversation, Message]: +def _runtime_records(session: Session) -> tuple[App, Conversation, Message]: + app = session.get(App, "app1") conversation = session.get(Conversation, "conv") message = session.get(Message, "msg") + assert app is not None assert conversation is not None assert message is not None - return conversation, message + session.expunge_all() + session.rollback() + return app, conversation, message class TestAgentChatAppRunnerRun: - def test_run_app_not_found( - self, - runner: AgentChatAppRunner, - sqlite_session: Session, - agent_generate_entity: AgentChatAppGenerateEntity, - agent_queue_manager: MessageBasedAppQueueManager, - ) -> None: - app = sqlite_session.get(App, "app1") - assert app is not None - sqlite_session.delete(app) - sqlite_session.commit() - conversation, message = _records(sqlite_session) - - with pytest.raises(ValueError, match="App not found"): - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) - def test_run_moderation_error_direct_output( self, runner: AgentChatAppRunner, @@ -73,15 +74,23 @@ def test_run_moderation_error_direct_output( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, ) -> None: def reject_input(_moderation: InputModeration, **kwargs: object) -> None: assert kwargs["query"] == "q" raise ModerationError("bad") monkeypatch.setattr(InputModeration, "check", reject_input) - conversation, message = _records(sqlite_session) + app, conversation, message = _runtime_records(sqlite_session) - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) _assert_direct_output(agent_queue_manager, "bad") @@ -92,6 +101,7 @@ def test_run_annotation_reply_short_circuits( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, ) -> None: annotation = MessageAnnotation( app_id="app1", @@ -101,21 +111,32 @@ def test_run_annotation_reply_short_circuits( ) sqlite_session.add(annotation) sqlite_session.commit() - sessions: list[Session] = [] + calls: list[dict[str, object]] = [] - def query_annotation(*, session: Session, **_kwargs: object) -> MessageAnnotation | None: - sessions.append(session) - return session.get(MessageAnnotation, annotation.id) + def query_annotation(**kwargs: object) -> MessageAnnotation: + calls.append(kwargs) + return annotation - monkeypatch.setattr(runner, "query_app_annotations_to_reply", query_annotation) - conversation, message = _records(sqlite_session) - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) + monkeypatch.setattr(annotation_replies, "query", query_annotation) + app, conversation, message = _runtime_records(sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) event_message = agent_queue_manager._q.get_nowait() assert event_message is not None assert isinstance(event_message.event, QueueAnnotationReplyEvent) assert event_message.event.message_annotation_id == annotation.id - assert sessions == [sqlite_session] + assert calls[0]["tenant_id"] == "tenant" + assert calls[0]["app_id"] == "app1" + assert calls[0]["message_id"] == "msg" + assert calls[0]["query"] == "q" + assert calls[0]["user_id"] == "user" _assert_direct_output(agent_queue_manager, "answer") def test_run_hosting_moderation_short_circuits( @@ -125,6 +146,7 @@ def test_run_hosting_moderation_short_circuits( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, ) -> None: def reject_hosted_input(**kwargs: object) -> bool: assert kwargs["tenant_id"] == "tenant" @@ -132,9 +154,16 @@ def reject_hosted_input(**kwargs: object) -> bool: return True monkeypatch.setattr("core.helper.moderation.check_moderation", reject_hosted_input) - conversation, message = _records(sqlite_session) + app, conversation, message = _runtime_records(sqlite_session) - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) _assert_direct_output( agent_queue_manager, @@ -148,12 +177,20 @@ def test_run_model_schema_missing( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, ) -> None: monkeypatch.setattr(PluginModelClient, "get_model_schema", lambda *_args, **_kwargs: None) - conversation, message = _records(sqlite_session) + app, conversation, message = _runtime_records(sqlite_session) with pytest.raises(ValueError, match="Model schema not found"): - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) @pytest.mark.parametrize( ("mode", "expected_runner"), @@ -172,6 +209,7 @@ def test_run_chain_of_thought_modes( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, ) -> None: agent_generate_entity.model_conf.mode = mode agent_generate_entity.model_conf.model_schema.model_properties[ModelPropertyKey.MODE] = mode @@ -183,12 +221,12 @@ def test_run_chain_of_thought_modes( original_run = CotAgentRunner.run def observe_run( - agent: CotAgentRunner, *, session: Session, message: Message, query: str, inputs: Mapping[str, str] + agent: CotAgentRunner, *, message: Message, query: str, inputs: Mapping[str, str] ) -> Generator[LLMResultChunk]: assert not sqlite_session.in_transaction() selected_runners.append(agent) events.append("agent-run") - return original_run(agent, session=session, message=message, query=query, inputs=inputs) + return original_run(agent, message=message, query=query, inputs=inputs) def dispatch( _client: PluginModelClient, **kwargs: object @@ -213,19 +251,19 @@ def dispatch( ), ) - def record_commit(_session: Session) -> None: - events.append("commit") - monkeypatch.setattr(CotAgentRunner, "run", observe_run) monkeypatch.setattr(PluginModelClient, "_request_with_plugin_daemon_response_stream", dispatch) - conversation, message = _records(sqlite_session) - event.listen(sqlite_session, "after_commit", record_commit) - try: - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) - finally: - event.remove(sqlite_session, "after_commit", record_commit) - - assert events == ["commit", "commit", "agent-run", "commit", "invoke"] + app, conversation, message = _runtime_records(sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) + + assert events == ["agent-run", "invoke"] assert len(selected_runners) == 1 assert type(selected_runners[0]) is expected_runner assert len(daemon_calls) == 1 @@ -251,12 +289,19 @@ def test_run_invalid_llm_mode_raises( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, ) -> None: - # Simulate corrupt provider metadata without replacing the runtime or validated app config. agent_generate_entity.model_conf.model_schema.model_properties[ModelPropertyKey.MODE] = "invalid" - conversation, message = _records(sqlite_session) + app, conversation, message = _runtime_records(sqlite_session) with pytest.raises(ValueError, match="Invalid LLM mode: invalid"): - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) @pytest.mark.parametrize("feature", [ModelFeature.TOOL_CALL, ModelFeature.MULTI_TOOL_CALL]) @pytest.mark.parametrize("stream_tool_call", [False, True]) @@ -268,6 +313,7 @@ def test_run_function_calling_strategy_selected_by_features( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, feature: ModelFeature, stream_tool_call: bool, ) -> None: @@ -278,11 +324,11 @@ def test_run_function_calling_strategy_selected_by_features( original_run = FunctionCallAgentRunner.run def run_agent( - agent: FunctionCallAgentRunner, *, session: Session, message: Message, query: str, **kwargs: object + agent: FunctionCallAgentRunner, *, message: Message, query: str, **kwargs: object ) -> Generator[LLMResultChunk]: - assert not session.in_transaction() + assert not sqlite_session.in_transaction() selected.append(agent) - return original_run(agent, session=session, message=message, query=query, **kwargs) + return original_run(agent, message=message, query=query, **kwargs) def dispatch(_client: PluginModelClient, **kwargs: object) -> Iterator[LLMResultChunk]: assert not sqlite_session.in_transaction() @@ -297,8 +343,15 @@ def dispatch(_client: PluginModelClient, **kwargs: object) -> Iterator[LLMResult monkeypatch.setattr(FunctionCallAgentRunner, "run", run_agent) monkeypatch.setattr(PluginModelClient, "_request_with_plugin_daemon_response_stream", dispatch) - conversation, message = _records(sqlite_session) - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) + app, conversation, message = _runtime_records(sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) app_config = agent_generate_entity.app_config assert isinstance(app_config, AgentChatAppConfig) @@ -321,33 +374,6 @@ def dispatch(_client: PluginModelClient, **kwargs: object) -> Iterator[LLMResult assert len(thoughts) == 1 assert thoughts[0].answer == "answer" - def test_run_conversation_not_found( - self, - runner: AgentChatAppRunner, - sqlite_session: Session, - agent_generate_entity: AgentChatAppGenerateEntity, - agent_queue_manager: MessageBasedAppQueueManager, - ) -> None: - conversation_record, message_record = _records(sqlite_session) - sqlite_session.delete(conversation_record) - sqlite_session.commit() - - with pytest.raises(ValueError, match="Conversation not found"): - runner.run(agent_generate_entity, agent_queue_manager, conversation_record, message_record, sqlite_session) - - def test_run_message_not_found( - self, - runner: AgentChatAppRunner, - sqlite_session: Session, - agent_generate_entity: AgentChatAppGenerateEntity, - agent_queue_manager: MessageBasedAppQueueManager, - ) -> None: - conversation, message = _records(sqlite_session) - sqlite_session.delete(message) - sqlite_session.commit() - with pytest.raises(ValueError, match="Message not found"): - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) - def test_run_invalid_agent_strategy_raises( self, runner: AgentChatAppRunner, @@ -355,13 +381,20 @@ def test_run_invalid_agent_strategy_raises( sqlite_session: Session, agent_generate_entity: AgentChatAppGenerateEntity, agent_queue_manager: MessageBasedAppQueueManager, + annotation_replies: AnnotationReplies, ) -> None: app_config = agent_generate_entity.app_config assert isinstance(app_config, AgentChatAppConfig) agent = app_config.agent assert isinstance(agent, AgentEntity) - # Fault injection after validation covers the defensive runtime branch for corrupted state. monkeypatch.setattr(agent, "strategy", "invalid") - conversation, message = _records(sqlite_session) + app, conversation, message = _runtime_records(sqlite_session) with pytest.raises(ValueError, match="Invalid agent strategy: invalid"): - runner.run(agent_generate_entity, agent_queue_manager, conversation, message, sqlite_session) + runner.run( + agent_generate_entity, + agent_queue_manager, + conversation, + message, + app_record=app, + annotations=annotation_replies, + ) diff --git a/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py index cf438c6824c8b2..8702bc8c2664ef 100644 --- a/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py @@ -1,22 +1,22 @@ from decimal import Decimal from types import SimpleNamespace -from unittest.mock import ANY, MagicMock, Mock, patch +from unittest.mock import Mock, patch import pytest -from sqlalchemy import event from sqlalchemy.orm import Session -import core.app.apps.chat.app_generator as generator_module -from core.app.apps.chat.app_generator import ChatAppGenerator -from core.app.apps.chat.app_runner import ChatAppRunner from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import QueueAnnotationReplyEvent from core.moderation.base import ModerationError +from extensions.application_services.retrieval import build_dataset_retrieval from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from models import Account from models.enums import ConversationFromSource from models.model import App, AppMode, AppModelConfig, Conversation, IconType, Message, MessageAnnotation +from services.app.generation.adapters.chat import ChatAppGenerator +from services.app.generation.adapters.chat_runner import ChatAppRunner +from services.errors.conversation import ConversationNotExistsError class DummyGenerateEntity: @@ -113,14 +113,15 @@ def _persist_records(session: Session) -> tuple[App, Conversation, Message]: return app, conversation, message -@pytest.fixture(autouse=True) -def _bind_db_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(generator_module.db, "session", sqlite_session) - - class TestChatAppGenerator: - def test_generate_requires_query(self, unbound_session: Session): - generator = ChatAppGenerator() + def test_generate_requires_query( + self, app_records, unbound_session: Session, *, annotation_replies, sqlite_session_factory + ): + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) with pytest.raises(ValueError, match="query is required"): generator.generate( session=unbound_session, @@ -131,8 +132,14 @@ def test_generate_requires_query(self, unbound_session: Session): streaming=False, ) - def test_generate_rejects_non_string_query(self, unbound_session: Session): - generator = ChatAppGenerator() + def test_generate_rejects_non_string_query( + self, app_records, unbound_session: Session, *, annotation_replies, sqlite_session_factory + ): + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) with pytest.raises(ValueError, match="query must be a string"): generator.generate( session=unbound_session, @@ -143,8 +150,14 @@ def test_generate_rejects_non_string_query(self, unbound_session: Session): streaming=False, ) - def test_generate_debugger_overrides_model_config(self, unbound_session: Session): - generator = ChatAppGenerator() + def test_generate_debugger_overrides_model_config( + self, app_records, unbound_session: Session, *, annotation_replies, sqlite_session_factory + ): + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) app_model = _app() user = _account() args = { @@ -157,26 +170,27 @@ def test_generate_debugger_overrides_model_config(self, unbound_session: Session with ( patch( - "core.app.apps.chat.app_generator.ConversationService.get_conversation", return_value=None + "repositories.app.generation_repository.AppGenerationRepository.conversation", return_value=None ) as get_conversation, - patch("core.app.apps.chat.app_generator.ChatAppConfigManager.config_validate", return_value={"x": 1}), + patch("services.app.generation.adapters.chat.ChatAppConfigManager.config_validate", return_value={"x": 1}), patch( - "core.app.apps.chat.app_generator.ChatAppConfigManager.get_app_config", + "services.app.generation.adapters.chat.ChatAppConfigManager.get_app_config", return_value=SimpleNamespace( variables=[], external_data_variables=[], app_model_config_dict={}, app_mode=AppMode.CHAT ), ), - patch("core.app.apps.chat.app_generator.ModelConfigConverter.convert", return_value=SimpleNamespace()), - patch("core.app.apps.chat.app_generator.FileUploadConfigManager.convert", return_value=None), - patch("core.app.apps.chat.app_generator.file_factory.build_from_mappings", return_value=[]), + patch("services.app.generation.adapters.chat.ModelConfigConverter.convert", return_value=SimpleNamespace()), + patch("services.app.generation.adapters.chat.FileUploadConfigManager.convert", return_value=None), + patch("services.app.generation.adapters.chat.file_factory.build_from_mappings", return_value=[]), patch( - "core.app.apps.chat.app_generator.ChatAppGenerateEntity", + "services.app.generation.adapters.chat.ChatAppGenerateEntity", Mock(side_effect=DummyGenerateEntity), ) as generate_entity, - patch("core.app.apps.chat.app_generator.TraceQueueManager", return_value=SimpleNamespace()), - patch("core.app.apps.chat.app_generator.MessageBasedAppQueueManager", DummyQueueManager), + patch("services.app.generation.adapters.chat.TraceQueueManager", return_value=SimpleNamespace()), + patch("services.app.generation.adapters.chat.MessageBasedAppQueueManager", DummyQueueManager), patch( - "core.app.apps.chat.app_generator.ChatAppGenerateResponseConverter.convert", return_value={"ok": True} + "services.app.generation.adapters.chat.ChatAppGenerateResponseConverter.convert", + return_value={"ok": True}, ), patch.object(ChatAppGenerator, "_get_app_model_config", return_value=AppModelConfig(app_id="app-1")), patch.object(ChatAppGenerator, "_prepare_user_inputs", return_value={}), @@ -186,8 +200,8 @@ def test_generate_debugger_overrides_model_config(self, unbound_session: Session return_value=(_conversation(), _message()), ), patch.object(ChatAppGenerator, "_handle_response", return_value={"response": True}), - patch("core.app.apps.chat.app_generator.copy_current_request_context", side_effect=lambda f: f), - patch("core.app.apps.chat.app_generator.threading.Thread") as mock_thread, + patch("services.app.generation.adapters.chat.copy_current_request_context", side_effect=lambda f: f), + patch("services.app.generation.adapters.chat.threading.Thread") as mock_thread, ): mock_thread.return_value.start.return_value = None result = generator.generate( @@ -195,11 +209,88 @@ def test_generate_debugger_overrides_model_config(self, unbound_session: Session ) assert result == {"ok": True} - assert get_conversation.call_args.kwargs["session"] is unbound_session + assert get_conversation.call_args.kwargs["app_id"] == "app-1" assert generate_entity.call_args.kwargs["extras"]["trace_session_id"] == "session-1" - def test_generate_uses_session_for_annotation_reply(self, unbound_session: Session): - generator = ChatAppGenerator() + def test_generate_releases_explicit_session_before_stream_consumption( + self, app_records, sqlite_session: Session, *, annotation_replies, sqlite_session_factory + ) -> None: + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) + app_model = _app() + sqlite_session.add(app_model) + sqlite_session.commit() + sqlite_session.expunge_all() + assert sqlite_session.get(App, app_model.id) is not None + assert sqlite_session.in_transaction() + + app_config = SimpleNamespace( + variables=[], + external_data_variables=[], + app_model_config_dict={}, + app_mode=AppMode.CHAT, + ) + events: list[str] = [] + + def initialize_records(*_args, **_kwargs): + assert not sqlite_session.in_transaction() + events.append("records") + return _conversation(), _message() + + def stream_response(**_kwargs): + assert not sqlite_session.in_transaction() + + def chunks(): + assert not sqlite_session.in_transaction() + events.append("first-chunk") + yield "chunk" + + return chunks() + + with ( + patch.object(ChatAppGenerator, "_get_app_model_config", return_value=AppModelConfig(app_id="app-1")), + patch("services.app.generation.adapters.chat.ChatAppConfigManager.config_validate", return_value={"x": 1}), + patch("services.app.generation.adapters.chat.ChatAppConfigManager.get_app_config", return_value=app_config), + patch("services.app.generation.adapters.chat.ModelConfigConverter.convert", return_value=SimpleNamespace()), + patch("services.app.generation.adapters.chat.FileUploadConfigManager.convert", return_value=None), + patch("services.app.generation.adapters.chat.ChatAppGenerateEntity", Mock(side_effect=DummyGenerateEntity)), + patch("services.app.generation.adapters.chat.TraceQueueManager", return_value=SimpleNamespace()), + patch("services.app.generation.adapters.chat.MessageBasedAppQueueManager", DummyQueueManager), + patch.object(ChatAppGenerator, "_prepare_user_inputs", return_value={}), + patch.object(ChatAppGenerator, "_init_generate_records", side_effect=initialize_records), + patch.object(ChatAppGenerator, "_handle_response", side_effect=stream_response), + patch( + "services.app.generation.adapters.chat.ChatAppGenerateResponseConverter.convert", + side_effect=lambda *, response, invoke_from: response, + ), + patch("services.app.generation.adapters.chat.copy_current_request_context", side_effect=lambda f: f), + patch("services.app.generation.adapters.chat.threading.Thread") as worker_thread, + ): + worker_thread.return_value.start.return_value = None + response = generator.generate( + app_model, + _account(), + {"query": "hi", "inputs": {}, "model_config": {"x": 1}}, + InvokeFrom.DEBUGGER, + streaming=True, + session=sqlite_session, + ) + + assert events == ["records"] + assert next(response) == "chunk" + assert events == ["records", "first-chunk"] + + def test_generate_uses_session_for_annotation_reply( + self, app_records, unbound_session: Session, *, annotation_replies, sqlite_session_factory + ): + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) app_model = _app() app_model_config = AppModelConfig(app_id="app-1") annotation_reply = {"enabled": False} @@ -209,12 +300,12 @@ def test_generate_uses_session_for_annotation_reply(self, unbound_session: Sessi patch.object(ChatAppGenerator, "_get_app_model_config", return_value=app_model_config), patch.object(AppModelConfig, "to_dict", return_value={}) as to_dict, patch( - "core.app.apps.chat.app_generator.load_annotation_reply_config", + "repositories.app.generation_repository.AppGenerationRepository.annotation_config", return_value=annotation_reply, ) as load_annotation_reply_config, - patch("core.app.apps.chat.app_generator.FileUploadConfigManager.convert", return_value=None), + patch("services.app.generation.adapters.chat.FileUploadConfigManager.convert", return_value=None), patch( - "core.app.apps.chat.app_generator.ChatAppConfigManager.get_app_config", + "services.app.generation.adapters.chat.ChatAppConfigManager.get_app_config", side_effect=RuntimeError("stop after app config"), ) as get_app_config, ): @@ -227,12 +318,18 @@ def test_generate_uses_session_for_annotation_reply(self, unbound_session: Sessi session=unbound_session, ) - load_annotation_reply_config.assert_called_once_with(unbound_session, "app-1") + load_annotation_reply_config.assert_called_once_with(tenant_id="tenant-1", app_id="app-1") to_dict.assert_called_once_with(annotation_reply=annotation_reply) assert get_app_config.call_args.kwargs["annotation_reply"] is annotation_reply - def test_generate_rejects_model_config_override_for_non_debugger(self, unbound_session: Session): - generator = ChatAppGenerator() + def test_generate_rejects_model_config_override_for_non_debugger( + self, app_records, unbound_session: Session, *, annotation_replies, sqlite_session_factory + ): + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) with pytest.raises(ValueError, match="Only in App debug mode can override model config"): with ( patch.object(ChatAppGenerator, "_get_app_model_config", return_value=AppModelConfig(app_id="app-1")), @@ -247,15 +344,18 @@ def test_generate_rejects_model_config_override_for_non_debugger(self, unbound_s streaming=False, ) - def test_generate_worker_handles_exceptions(self): - generator = ChatAppGenerator() + def test_generate_worker_handles_exceptions(self, app_records, *, annotation_replies, sqlite_session_factory): + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) queue_manager = DummyQueueManager() - entity = DummyGenerateEntity(task_id="t1", user_id="u1") + entity = DummyGenerateEntity(task_id="t1", user_id="u1", app_config=Mock(app_id="app-1", tenant_id="tenant-1")) with ( - patch.object(ChatAppGenerator, "_get_conversation", return_value=_conversation()), - patch.object(ChatAppGenerator, "_get_message", return_value=_message()), - patch("core.app.apps.chat.app_generator.ChatAppRunner.run", side_effect=InvokeAuthorizationError()), + patch.object(app_records, "load", return_value=(_app(), _conversation(), _message())), + patch("services.app.generation.adapters.chat.ChatAppRunner.run", side_effect=InvokeAuthorizationError()), ): generator._generate_worker( flask_app=Mock(app_context=Mock(return_value=Mock(__enter__=Mock(), __exit__=Mock()))), @@ -268,9 +368,8 @@ def test_generate_worker_handles_exceptions(self): assert queue_manager.published with ( - patch.object(ChatAppGenerator, "_get_conversation", return_value=_conversation()), - patch.object(ChatAppGenerator, "_get_message", return_value=_message()), - patch("core.app.apps.chat.app_generator.ChatAppRunner.run", side_effect=GenerateTaskStoppedError()), + patch.object(app_records, "load", return_value=(_app(), _conversation(), _message())), + patch("services.app.generation.adapters.chat.ChatAppRunner.run", side_effect=GenerateTaskStoppedError()), ): generator._generate_worker( flask_app=Mock(app_context=Mock(return_value=Mock(__enter__=Mock(), __exit__=Mock()))), @@ -280,38 +379,42 @@ def test_generate_worker_handles_exceptions(self): message_id="m1", ) - -class TestChatAppRunner: - def test_run_raises_when_app_missing(self, sqlite_session: Session): - runner = ChatAppRunner() - app_config = SimpleNamespace( - app_id="app-1", tenant_id="tenant-1", prompt_template=None, external_data_variables=[] + def test_generate_worker_reports_missing_app( + self, app, app_records, *, annotation_replies, sqlite_session_factory + ) -> None: + generator = ChatAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, ) - app_generate_entity = DummyGenerateEntity( - app_config=app_config, - model_conf=SimpleNamespace(provider_model_bundle=None, model=None, parameters={}, app_model_config_dict={}), - inputs={}, - query="hi", - files=[], - file_upload_config=None, - conversation_id=None, - stream=False, - user_id="user-1", - invoke_from=InvokeFrom.SERVICE_API, + queue_manager = DummyQueueManager() + entity = DummyGenerateEntity( + task_id="t1", + user_id="u1", + app_config=SimpleNamespace(app_id="missing-app", tenant_id="tenant-1"), ) - with pytest.raises(ValueError, match="App not found"): - runner.run( - app_generate_entity, - DummyQueueManager(), - _conversation(), - _message(), - sqlite_session, - ) + generator._generate_worker( + flask_app=app, + application_generate_entity=entity, + queue_manager=queue_manager, + conversation_id="missing-conversation", + message_id="missing-message", + ) + + assert len(queue_manager.published) == 1 + assert isinstance(queue_manager.published[0][0], ConversationNotExistsError) - def test_run_moderation_error_direct_output(self, sqlite_session: Session): + +class TestChatAppRunner: + def test_run_moderation_error_direct_output( + self, sqlite_session: Session, annotation_replies, *, app_records, sqlite_session_factory + ): _, conversation, message = _persist_records(sqlite_session) - runner = ChatAppRunner() + runner = ChatAppRunner( + retrieval=build_dataset_retrieval(sqlite_session_factory), + records=app_records, + ) app_config = SimpleNamespace( app_id="app-1", tenant_id="tenant-1", @@ -344,14 +447,20 @@ def test_run_moderation_error_direct_output(self, sqlite_session: Session): DummyQueueManager(), conversation, message, - sqlite_session, + app_record=sqlite_session.get(App, "app-1"), + annotations=annotation_replies, ) mock_direct.assert_called_once() - def test_run_annotation_reply_short_circuits(self, sqlite_session: Session): + def test_run_annotation_reply_short_circuits( + self, sqlite_session: Session, annotation_replies, *, app_records, sqlite_session_factory + ): _, conversation, message = _persist_records(sqlite_session) - runner = ChatAppRunner() + runner = ChatAppRunner( + retrieval=build_dataset_retrieval(sqlite_session_factory), + records=app_records, + ) app_config = SimpleNamespace( app_id="app-1", tenant_id="tenant-1", @@ -384,19 +493,31 @@ def test_run_annotation_reply_short_circuits(self, sqlite_session: Session): with ( patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")), - patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=annotation) as annotation_query, + patch.object(annotation_replies, "query", return_value=annotation) as annotation_query, patch.object(ChatAppRunner, "direct_output") as mock_direct, ): queue_manager = DummyQueueManager() - runner.run(app_generate_entity, queue_manager, conversation, message, sqlite_session) + runner.run( + app_generate_entity, + queue_manager, + conversation, + message, + app_record=sqlite_session.get(App, "app-1"), + annotations=annotation_replies, + ) assert any(isinstance(item[0], QueueAnnotationReplyEvent) for item in queue_manager.published) - assert annotation_query.call_args.kwargs["session"] is sqlite_session + assert annotation_query.call_args.kwargs["app_id"] == "app-1" mock_direct.assert_called_once() - def test_run_returns_when_hosting_moderation_blocks(self, sqlite_session: Session): + def test_run_returns_when_hosting_moderation_blocks( + self, sqlite_session: Session, annotation_replies, *, app_records, sqlite_session_factory + ): _, conversation, message = _persist_records(sqlite_session) - runner = ChatAppRunner() + runner = ChatAppRunner( + retrieval=build_dataset_retrieval(sqlite_session_factory), + records=app_records, + ) app_config = SimpleNamespace( app_id="app-1", tenant_id="tenant-1", @@ -421,7 +542,7 @@ def test_run_returns_when_hosting_moderation_blocks(self, sqlite_session: Sessio with ( patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")), - patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=None), + patch.object(annotation_replies, "query", return_value=None), patch.object(ChatAppRunner, "check_hosting_moderation", return_value=True), ): runner.run( @@ -429,75 +550,6 @@ def test_run_returns_when_hosting_moderation_blocks(self, sqlite_session: Sessio DummyQueueManager(), conversation, message, - sqlite_session, + app_record=sqlite_session.get(App, "app-1"), + annotations=annotation_replies, ) - - def test_run_closes_explicit_session_before_stream_consumption(self, sqlite_session: Session): - _, conversation, message = _persist_records(sqlite_session) - runner = ChatAppRunner() - app_config = SimpleNamespace( - app_id="app-1", - tenant_id="tenant-1", - app_mode=AppMode.CHAT, - prompt_template=None, - external_data_variables=[], - dataset=None, - additional_features=None, - ) - app_generate_entity = DummyGenerateEntity( - app_config=app_config, - model_conf=SimpleNamespace(provider_model_bundle=None, model="model-1", parameters={}), - inputs={}, - query="hi", - files=[], - file_upload_config=None, - conversation_id=None, - stream=True, - user_id="user-1", - invoke_from=InvokeFrom.SERVICE_API, - ) - - events = [] - queue_manager = DummyQueueManager() - model_instance = MagicMock() - - def record_commit(_session: Session) -> None: - events.append("commit") - - def invoke_stream(): - events.append("first-chunk") - yield "chunk" - - def invoke_llm(**kwargs): - events.append("invoke") - return invoke_stream() - - event.listen(sqlite_session, "after_commit", record_commit) - try: - with ( - patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), - patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")), - patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=None), - patch.object(ChatAppRunner, "check_hosting_moderation", return_value=False), - patch.object(ChatAppRunner, "recalc_llm_max_tokens"), - patch.object( - ChatAppRunner, - "_handle_invoke_result", - side_effect=lambda invoke_result, **kwargs: list(invoke_result), - ) as mock_handle, - patch("core.app.apps.chat.app_runner.ModelInstance", return_value=model_instance), - ): - model_instance.invoke_llm.side_effect = invoke_llm - runner.run(app_generate_entity, queue_manager, conversation, message, sqlite_session) - finally: - event.remove(sqlite_session, "after_commit", record_commit) - - assert events == ["commit", "commit", "invoke", "first-chunk"] - mock_handle.assert_called_once_with( - invoke_result=ANY, - queue_manager=queue_manager, - stream=True, - message_id="m1", - user_id="user-1", - tenant_id="tenant-1", - ) diff --git a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py index 9d13867b1a5cd7..9c18893fe62b46 100644 --- a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py +++ b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py @@ -3,7 +3,6 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import UTC, datetime, timedelta -from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock @@ -11,9 +10,6 @@ from sqlalchemy.orm import Session from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.common import workflow_response_converter -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.apps.workflow.app_runner import WorkflowAppRunner from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.app.entities.queue_entities import ( NodeExecutionSnapshot, @@ -49,7 +45,10 @@ from models.human_input import HumanInputForm, HumanInputFormRecipient from models.human_input_entities import ParagraphInputConfig, SelectInputConfig, StringListSource, UserActionConfig from models.model import AppMode -from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowType +from models.workflow import WorkflowNodeExecutionModel +from services.workflow.execution.adapters import response_converter as workflow_response_converter +from services.workflow.execution.adapters.events import WorkflowEventPublisher +from services.workflow.execution.adapters.response_converter import WorkflowResponseConverter class TestWorkflowResponseConverterCases: @@ -306,7 +305,9 @@ def test_fetch_files_from_variable_value_with_complex_nested_scenario(self): assert result[1]["id"] == "complex_obj" -def _workflow_response_converter_human_input_build_converter() -> WorkflowResponseConverter: +def _workflow_response_converter_human_input_build_converter( + *, workflow_contexts, tool_providers +) -> WorkflowResponseConverter: system_variables = build_system_variables( files=[], user_id="user-1", @@ -331,16 +332,22 @@ def _workflow_response_converter_human_input_build_converter() -> WorkflowRespon account = _workflow_response_converter_human_input_Account(name="tester", email="tester@example.com") account.id = "acc-1" return WorkflowResponseConverter( + contexts=workflow_contexts, application_generate_entity=app_entity, user=account, system_variables=system_variables, + tool_providers=tool_providers, ) class TestWorkflowResponseConverterHumanInputCases: @staticmethod - def test_human_input_form_filled_stream_response_contains_rendered_content() -> None: - converter = _workflow_response_converter_human_input_build_converter() + def test_human_input_form_filled_stream_response_contains_rendered_content( + *, workflow_contexts, tool_providers + ) -> None: + converter = _workflow_response_converter_human_input_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", @@ -368,8 +375,12 @@ def test_human_input_form_filled_stream_response_contains_rendered_content() -> assert resp.data.action_id == "Approve" @staticmethod - def test_human_input_form_filled_stream_response_serializes_submitted_data() -> None: - converter = _workflow_response_converter_human_input_build_converter() + def test_human_input_form_filled_stream_response_serializes_submitted_data( + *, workflow_contexts, tool_providers + ) -> None: + converter = _workflow_response_converter_human_input_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", @@ -399,8 +410,12 @@ def test_human_input_form_filled_stream_response_serializes_submitted_data() -> } @staticmethod - def test_human_input_form_timeout_stream_response_contains_timeout_metadata() -> None: - converter = _workflow_response_converter_human_input_build_converter() + def test_human_input_form_timeout_stream_response_contains_timeout_metadata( + *, workflow_contexts, tool_providers + ) -> None: + converter = _workflow_response_converter_human_input_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", @@ -425,7 +440,9 @@ def test_human_input_form_timeout_stream_response_contains_timeout_metadata() -> assert resp.data.expiration_time == 1735689600 -def _workflow_response_converter_resumption_build_converter() -> WorkflowResponseConverter: +def _workflow_response_converter_resumption_build_converter( + *, workflow_contexts, tool_providers +) -> WorkflowResponseConverter: """Construct a minimal WorkflowResponseConverter for testing.""" system_variables = build_system_variables( files=[], @@ -451,16 +468,20 @@ def _workflow_response_converter_resumption_build_converter() -> WorkflowRespons account = _workflow_response_converter_resumption_Account(name="tester", email="tester@example.com") account.id = "acc-1" return WorkflowResponseConverter( + contexts=workflow_contexts, application_generate_entity=app_entity, user=account, system_variables=system_variables, + tool_providers=tool_providers, ) class TestWorkflowResponseConverterResumptionCases: @staticmethod - def test_workflow_start_stream_response_carries_resumption_reason() -> None: - converter = _workflow_response_converter_resumption_build_converter() + def test_workflow_start_stream_response_carries_resumption_reason(*, workflow_contexts, tool_providers) -> None: + converter = _workflow_response_converter_resumption_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) resp = converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", @@ -470,8 +491,10 @@ def test_workflow_start_stream_response_carries_resumption_reason() -> None: assert resp.data.reason is WorkflowStartReason.RESUMPTION @staticmethod - def test_workflow_start_stream_response_carries_initial_reason() -> None: - converter = _workflow_response_converter_resumption_build_converter() + def test_workflow_start_stream_response_carries_initial_reason(*, workflow_contexts, tool_providers) -> None: + converter = _workflow_response_converter_resumption_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) resp = converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", @@ -481,8 +504,12 @@ def test_workflow_start_stream_response_carries_initial_reason() -> None: assert resp.data.reason is WorkflowStartReason.INITIAL @staticmethod - def test_node_start_preserves_the_allocated_index_when_events_arrive_out_of_order(): - converter = _workflow_response_converter_resumption_build_converter() + def test_node_start_preserves_the_allocated_index_when_events_arrive_out_of_order( + *, workflow_contexts, tool_providers + ): + converter = _workflow_response_converter_resumption_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL ) @@ -505,9 +532,14 @@ def test_node_start_preserves_the_allocated_index_when_events_arrive_out_of_orde @staticmethod @pytest.mark.parametrize("provider_type", ["workflow", "builtin"]) - def test_workflow_tool_detail_hint_in_live_and_saved_traces(monkeypatch, provider_type): + def test_workflow_tool_detail_hint_in_live_and_saved_traces( + monkeypatch, provider_type, *, workflow_contexts, tool_providers + ): monkeypatch.setattr(workflow_response_converter.ToolManager, "get_tool_icon", lambda **_kwargs: "icon") - converter = _workflow_response_converter_resumption_build_converter() + monkeypatch.setattr("core.tools.tool_manager.ToolManager.get_tool_icon", lambda **_kwargs: "icon") + converter = _workflow_response_converter_resumption_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL ) @@ -533,9 +565,11 @@ def test_workflow_tool_detail_hint_in_live_and_saved_traces(monkeypatch, provide @staticmethod @pytest.mark.parametrize("next_node_already_persisted", [False, True]) - def test_resumed_node_metadata_and_sequence(next_node_already_persisted): + def test_resumed_node_metadata_and_sequence(next_node_already_persisted, *, workflow_contexts, tool_providers): started_at = datetime(2026, 9, 7) - converter = _workflow_response_converter_resumption_build_converter() + converter = _workflow_response_converter_resumption_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL ) @@ -570,7 +604,9 @@ def node_start(execution_id, title, node_type, index): QueueWorkflowStartedEvent(node_execution_snapshots=tuple(snapshots)).model_dump_json() ) - converter = _workflow_response_converter_resumption_build_converter() + converter = _workflow_response_converter_resumption_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", @@ -636,7 +672,7 @@ def create_generate_entity(self) -> WorkflowAppGenerateEntity: stream=True, ) - def create_workflow_response_converter(self) -> WorkflowResponseConverter: + def create_workflow_response_converter(self, *, workflow_contexts, tool_providers) -> WorkflowResponseConverter: """Create a WorkflowResponseConverter for testing.""" entity = self.create_generate_entity() @@ -648,9 +684,11 @@ def create_workflow_response_converter(self) -> WorkflowResponseConverter: system_variables = build_system_variables(workflow_id="wf-id", workflow_execution_id="initial-run-id") return WorkflowResponseConverter( + contexts=workflow_contexts, application_generate_entity=entity, user=mock_user, system_variables=system_variables, + tool_providers=tool_providers, ) def create_node_started_event(self, *, node_execution_id: str | None = None) -> QueueNodeStartedEvent: @@ -711,9 +749,11 @@ def create_node_retry_event( in_loop_id=None, ) - def test_workflow_node_finish_response_uses_truncated_process_data(self): + def test_workflow_node_finish_response_uses_truncated_process_data(self, *, workflow_contexts, tool_providers): """Test that node finish response uses get_response_process_data().""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) original_data = {"large_field": "x" * 10000, "metadata": "info"} truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"} @@ -753,9 +793,11 @@ def fake_truncate(mapping): assert response.data.process_data != original_data assert response.data.process_data_truncated is True - def test_workflow_node_finish_response_without_truncation(self): + def test_workflow_node_finish_response_without_truncation(self, *, workflow_contexts, tool_providers): """Test node finish response when no truncation is applied.""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) original_data = {"small": "data"} @@ -791,9 +833,11 @@ def fake_truncate(mapping): assert response.data.process_data == original_data assert response.data.process_data_truncated is False - def test_workflow_node_finish_response_with_none_process_data(self): + def test_workflow_node_finish_response_with_none_process_data(self, *, workflow_contexts, tool_providers): """Test node finish response when process_data is None.""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="bootstrap", @@ -830,15 +874,20 @@ def fake_truncate(mapping): def test_workflow_node_finish_response_prefers_event_finished_at( self, monkeypatch: pytest.MonkeyPatch, + *, + workflow_contexts, + tool_providers, ) -> None: """Finished timestamps should come from the event, not delayed queue processing time.""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) start_at = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC).replace(tzinfo=None) finished_at = datetime(2024, 1, 1, 0, 0, 2, tzinfo=UTC).replace(tzinfo=None) delayed_processing_time = datetime(2024, 1, 1, 0, 0, 10, tzinfo=UTC).replace(tzinfo=None) monkeypatch.setattr( - "core.app.apps.common.workflow_response_converter.naive_utc_now", + "services.workflow.execution.adapters.response_converter.naive_utc_now", lambda: delayed_processing_time, ) converter.workflow_start_to_stream_response( @@ -871,9 +920,11 @@ def test_workflow_node_finish_response_prefers_event_finished_at( assert response.data.elapsed_time == 2.0 assert response.data.finished_at == int(finished_at.timestamp()) - def test_workflow_node_retry_response_uses_truncated_process_data(self): + def test_workflow_node_retry_response_uses_truncated_process_data(self, *, workflow_contexts, tool_providers): """Test that node retry response uses get_response_process_data().""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) original_data = {"large_field": "x" * 10000, "metadata": "info"} truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"} @@ -913,9 +964,11 @@ def fake_truncate(mapping): assert response.data.process_data != original_data assert response.data.process_data_truncated is True - def test_workflow_node_retry_response_without_truncation(self): + def test_workflow_node_retry_response_without_truncation(self, *, workflow_contexts, tool_providers): """Test node retry response when no truncation is applied.""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) original_data = {"small": "data"} @@ -950,9 +1003,11 @@ def fake_truncate(mapping): assert response.data.process_data == original_data assert response.data.process_data_truncated is False - def test_iteration_and_loop_nodes_return_none(self): + def test_iteration_and_loop_nodes_return_none(self, *, workflow_contexts, tool_providers): """Test that iteration and loop nodes return None (no streaming events).""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) iteration_event = QueueNodeSucceededEvent( node_id="iteration-node", @@ -980,9 +1035,11 @@ def test_iteration_and_loop_nodes_return_none(self): ) assert response is None - def test_finish_without_start_raises(self): + def test_finish_without_start_raises(self, *, workflow_contexts, tool_providers): """Ensure finish responses require a prior workflow start.""" - converter = self.create_workflow_response_converter() + converter = self.create_workflow_response_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) event = self.create_node_succeeded_event( node_execution_id=str(uuid.uuid4()), process_data={}, @@ -1037,16 +1094,20 @@ def create_test_system_variables(self): """Create test system variables.""" return build_system_variables() - def create_test_converter(self, invoke_from: InvokeFrom) -> WorkflowResponseConverter: + def create_test_converter( + self, invoke_from: InvokeFrom, *, workflow_contexts, tool_providers + ) -> WorkflowResponseConverter: """Create WorkflowResponseConverter with specified invoke_from.""" entity = self.create_test_app_generate_entity(invoke_from) user = self.create_test_user() system_variables = self.create_test_system_variables() converter = WorkflowResponseConverter( + contexts=workflow_contexts, application_generate_entity=entity, user=user, system_variables=system_variables, + tool_providers=tool_providers, ) # ensure `workflow_run_id` is set. converter.workflow_start_to_stream_response( @@ -1093,9 +1154,13 @@ def create_test_converter(self, invoke_from: InvokeFrom) -> WorkflowResponseConv ], ids=lambda x: x.name, ) - def test_truncator_selection_based_on_invoke_from(self, test_case: TestCase): + def test_truncator_selection_based_on_invoke_from( + self, test_case: TestCase, *, workflow_contexts, tool_providers + ): """Test that the correct truncator is selected based on invoke_from.""" - converter = self.create_test_converter(test_case.invoke_from) + converter = self.create_test_converter( + test_case.invoke_from, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Test truncation behavior instead of checking private attribute @@ -1136,9 +1201,11 @@ def test_truncator_selection_based_on_invoke_from(self, test_case: TestCase): assert not response.data.process_data_truncated assert not response.data.outputs_truncated - def test_service_api_truncator_no_op_mapping(self): + def test_service_api_truncator_no_op_mapping(self, *, workflow_contexts, tool_providers): """Test that Service API truncator doesn't truncate variable mappings.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) + converter = self.create_test_converter( + InvokeFrom.SERVICE_API, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Create a test event with large data large_value: dict[str, Any] = { @@ -1176,9 +1243,11 @@ def test_service_api_truncator_no_op_mapping(self): assert data.process_data_truncated is False assert data.outputs_truncated is False - def test_web_app_truncator_works_normally(self): + def test_web_app_truncator_works_normally(self, *, workflow_contexts, tool_providers): """Test that web app truncator still works normally.""" - converter = self.create_test_converter(InvokeFrom.WEB_APP) + converter = self.create_test_converter( + InvokeFrom.WEB_APP, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Create a test event with large data large_value = { @@ -1276,9 +1345,13 @@ def _create_event_by_type( QueueEvent.NODE_EXCEPTION, ], ) - def test_service_api_node_finish_event_no_truncation(self, event_type: QueueEvent): + def test_service_api_node_finish_event_no_truncation( + self, event_type: QueueEvent, *, workflow_contexts, tool_providers + ): """Test that Service API doesn't truncate node finish events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) + converter = self.create_test_converter( + InvokeFrom.SERVICE_API, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Create test event with large data large_inputs = {"input1": "x" * 5000, "input2": list(range(2000))} large_process_data = {"process1": "y" * 5000, "process2": {"nested": ["z"] * 2000}} @@ -1306,9 +1379,11 @@ def test_service_api_node_finish_event_no_truncation(self, event_type: QueueEven assert not response.data.process_data_truncated assert not response.data.outputs_truncated - def test_service_api_node_retry_event_no_truncation(self): + def test_service_api_node_retry_event_no_truncation(self, *, workflow_contexts, tool_providers): """Test that Service API doesn't truncate node retry events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) + converter = self.create_test_converter( + InvokeFrom.SERVICE_API, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Create test event with large data large_inputs = {"retry_input": "x" * 5000} @@ -1367,9 +1442,11 @@ def test_service_api_node_retry_event_no_truncation(self): assert not response.data.process_data_truncated assert not response.data.outputs_truncated - def test_service_api_iteration_events_no_truncation(self): + def test_service_api_iteration_events_no_truncation(self, *, workflow_contexts, tool_providers): """Test that Service API doesn't truncate iteration events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) + converter = self.create_test_converter( + InvokeFrom.SERVICE_API, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Test iteration start event large_value = {"iteration_input": ["x"] * 2000} @@ -1394,9 +1471,11 @@ def test_service_api_iteration_events_no_truncation(self): assert response.data.inputs == large_value assert not response.data.inputs_truncated - def test_service_api_loop_events_no_truncation(self): + def test_service_api_loop_events_no_truncation(self, *, workflow_contexts, tool_providers): """Test that Service API doesn't truncate loop events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) + converter = self.create_test_converter( + InvokeFrom.SERVICE_API, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Test loop start event large_inputs = {"loop_input": ["x"] * 2000} @@ -1421,9 +1500,11 @@ def test_service_api_loop_events_no_truncation(self): assert response.data.inputs == large_inputs assert not response.data.inputs_truncated - def test_web_app_node_finish_event_truncation_works(self): + def test_web_app_node_finish_event_truncation_works(self, *, workflow_contexts, tool_providers): """Test that web app still truncates node finish events.""" - converter = self.create_test_converter(InvokeFrom.WEB_APP) + converter = self.create_test_converter( + InvokeFrom.WEB_APP, workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) # Create test event with large data that should be truncated large_inputs = {"input1": ["x"] * 2000} @@ -1464,12 +1545,11 @@ def test_web_app_node_finish_event_truncation_works(self): assert response.data.outputs_truncated -class _RecordingWorkflowAppRunner(WorkflowAppRunner): - def __init__(self, **kwargs): - super().__init__(**kwargs) +class _RecordingQueueManager: + def __init__(self) -> None: self.published_events = [] - def _publish_event(self, event): + def publish(self, event, _pub_from) -> None: self.published_events.append(event) @@ -1506,45 +1586,15 @@ def _persist_human_input_form( return expiration_time -def _build_runner(): - app_entity = SimpleNamespace( - app_config=SimpleNamespace(app_id="app-id"), - inputs={}, - files=[], - invoke_from=InvokeFrom.SERVICE_API, - single_iteration_run=None, - single_loop_run=None, - workflow_execution_id="run-id", - user_id="user-id", - ) - workflow = Workflow.new( - tenant_id="tenant-id", - app_id="app-id", - type=WorkflowType.WORKFLOW, - version=Workflow.VERSION_DRAFT, - graph=json.dumps({}), - features="{}", - created_by="account-id", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - workflow.id = "workflow-id" - queue_manager = SimpleNamespace(publish=lambda event, pub_from: None) - return _RecordingWorkflowAppRunner( - application_generate_entity=app_entity, - queue_manager=queue_manager, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id="sys-user", - root_node_id=None, - graph_engine_layers=(), - graph_runtime_state=None, - ) +def _build_event_publisher() -> tuple[WorkflowEventPublisher, _RecordingQueueManager, MagicMock]: + queue_manager = _RecordingQueueManager() + notify_pause = MagicMock() + publisher = WorkflowEventPublisher(queue_manager, notify_pause=notify_pause) # type: ignore[arg-type] + return publisher, queue_manager, notify_pause def _workflow_pause_events_build_converter( - *, invoke_from: InvokeFrom = InvokeFrom.SERVICE_API + *, workflow_contexts, tool_providers, invoke_from: InvokeFrom = InvokeFrom.SERVICE_API ) -> WorkflowResponseConverter: application_generate_entity = WorkflowAppGenerateEntity( user_from=UserFrom.END_USER, @@ -1568,18 +1618,27 @@ def _workflow_pause_events_build_converter( user = _workflow_pause_events_Account(name="Tester", email="tester@example.com") user.id = "account-id" return WorkflowResponseConverter( + contexts=workflow_contexts, application_generate_entity=application_generate_entity, user=user, system_variables=system_variables, + tool_providers=tool_providers, ) def _build_paused_human_input_response( session: Session, recipients: list[tuple[RecipientType, str]], + *, + workflow_contexts, + tool_providers, ): """Drive the live OPENAPI pause path with persisted forms and recipients.""" - converter = _workflow_pause_events_build_converter(invoke_from=InvokeFrom.OPENAPI) + converter = _workflow_pause_events_build_converter( + workflow_contexts=workflow_contexts, + tool_providers=tool_providers, + invoke_from=InvokeFrom.OPENAPI, + ) converter.workflow_start_to_stream_response( task_id="task", workflow_run_id="run-id", @@ -1615,20 +1674,22 @@ def _build_paused_human_input_response( class TestWorkflowPauseEventsCases: @staticmethod - @pytest.mark.parametrize("source_node_id", ["node-human", "nested-human"]) - def test_graph_run_paused_event_uses_projected_nodes_for_highlighting( - monkeypatch: pytest.MonkeyPatch, source_node_id: str - ): - from tests.unit_tests.core.app.apps.test_workflow_app_runner import _workflow_entry + @pytest.fixture + def sqlite_pause_session(sqlite_session: Session) -> Session: + """Provide the shared SQLite session used by the injected context repository.""" + return sqlite_session - runner = _build_runner() + @staticmethod + @pytest.mark.parametrize("source_node_id", ["node-human", "nested-human"]) + def test_graph_run_paused_event_uses_projected_nodes_for_highlighting(source_node_id: str): + publisher, queue_manager, notify_pause = _build_event_publisher() graph_reason = HitlRequired( session_id="form-1", node_id=source_node_id, node_title="Human Step", ) event = GraphRunPausedEvent(reasons=[graph_reason], outputs={"foo": "bar"}) - workflow_entry = _workflow_entry(runner) + workflow_entry = MagicMock() enriched_reason = HumanInputRequired( form_id="form-1", @@ -1638,12 +1699,11 @@ def test_graph_run_paused_event_uses_projected_nodes_for_highlighting( node_id="node-human", node_title="Human Step", ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_form_delivery_task", MagicMock()) - - runner.handle_event(workflow_entry, event, pause_reasons=[enriched_reason]) + publisher.publish(workflow_entry, event, pause_reasons=[enriched_reason]) - assert len(runner.published_events) == 1 - queue_event = runner.published_events[0] + notify_pause.assert_called_once_with([enriched_reason]) + assert len(queue_manager.published_events) == 1 + queue_event = queue_manager.published_events[0] assert isinstance(queue_event, QueueWorkflowPausedEvent) assert queue_event.reasons == [enriched_reason] assert queue_event.outputs == {"foo": "bar"} @@ -1656,8 +1716,12 @@ def test_graph_run_paused_event_uses_projected_nodes_for_highlighting( [(HumanInputForm, HumanInputFormRecipient)], indirect=True, ) - def test_queue_workflow_paused_event_to_stream_responses(sqlite_pause_session: Session): - converter = _workflow_pause_events_build_converter() + def test_queue_workflow_paused_event_to_stream_responses( + sqlite_pause_session: Session, *, workflow_contexts, tool_providers + ): + converter = _workflow_pause_events_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task", workflow_run_id="run-id", @@ -1719,13 +1783,17 @@ def test_queue_workflow_paused_event_to_stream_responses(sqlite_pause_session: S [(HumanInputForm, HumanInputFormRecipient)], indirect=True, ) - def test_openapi_pause_without_web_app_recipient_emits_approval_channels(sqlite_pause_session: Session): + def test_openapi_pause_without_web_app_recipient_emits_approval_channels( + sqlite_pause_session: Session, *, workflow_contexts, tool_providers + ): responses = _build_paused_human_input_response( sqlite_pause_session, recipients=[ (RecipientType.EMAIL_MEMBER, "email-token"), (RecipientType.BACKSTAGE, "backstage-token"), ], + workflow_contexts=workflow_contexts, + tool_providers=tool_providers, ) hi_resp = responses[0] @@ -1741,13 +1809,17 @@ def test_openapi_pause_without_web_app_recipient_emits_approval_channels(sqlite_ [(HumanInputForm, HumanInputFormRecipient)], indirect=True, ) - def test_openapi_pause_with_web_app_recipient_sets_token_and_channels(sqlite_pause_session: Session): + def test_openapi_pause_with_web_app_recipient_sets_token_and_channels( + sqlite_pause_session: Session, *, workflow_contexts, tool_providers + ): responses = _build_paused_human_input_response( sqlite_pause_session, recipients=[ (RecipientType.STANDALONE_WEB_APP, "web-app-token"), (RecipientType.BACKSTAGE, "backstage-token"), ], + workflow_contexts=workflow_contexts, + tool_providers=tool_providers, ) hi_resp = responses[0] @@ -1763,8 +1835,12 @@ def test_openapi_pause_with_web_app_recipient_sets_token_and_channels(sqlite_pau [(HumanInputForm, HumanInputFormRecipient)], indirect=True, ) - def test_queue_workflow_paused_event_resolves_variable_select_options(sqlite_pause_session: Session): - converter = _workflow_pause_events_build_converter() + def test_queue_workflow_paused_event_resolves_variable_select_options( + sqlite_pause_session: Session, *, workflow_contexts, tool_providers + ): + converter = _workflow_pause_events_build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) converter.workflow_start_to_stream_response( task_id="task", workflow_run_id="run-id", @@ -1815,7 +1891,9 @@ def test_queue_workflow_paused_event_resolves_variable_select_options(sqlite_pau @pytest.mark.parametrize("sqlite_session", [(HumanInputForm, HumanInputFormRecipient)], indirect=True) -def test_mixed_human_input_versions_preserve_order_and_wire_payloads(sqlite_pause_session: Session): +def test_mixed_human_input_versions_preserve_order_and_wire_payloads( + sqlite_session: Session, *, workflow_contexts, tool_providers +) -> None: from pydantic import SecretStr from core.human_input_v2.resolved_form import ResolvedForm, SelectInput @@ -1826,7 +1904,7 @@ def test_mixed_human_input_versions_preserve_order_and_wire_payloads(sqlite_paus from graphon.entities.pause_reason import SchedulingPause from repositories.human_input_v2.form_repository import Form - _persist_human_input_form(sqlite_pause_session) + _persist_human_input_form(sqlite_session) now = datetime(2024, 1, 1) form = Form( id="form-2", @@ -1874,7 +1952,10 @@ def test_mixed_human_input_versions_preserve_order_and_wire_payloads(sqlite_paus ) state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0) state.variable_pool.add(["upstream", "options"], ["One", "Two"]) - converter = _workflow_pause_events_build_converter() + converter = _workflow_pause_events_build_converter( + workflow_contexts=workflow_contexts, + tool_providers=tool_providers, + ) converter.workflow_start_to_stream_response( task_id="task", workflow_run_id="run-id", @@ -1898,10 +1979,3 @@ def test_mixed_human_input_versions_preserve_order_and_wire_payloads(sqlite_paus assert v1["form_version"] == "1" assert reasons[0]["form_version"] == "2" assert reasons[2]["form_version"] == "1" - - -@pytest.fixture -def sqlite_pause_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Bind pause-response queries to the shared SQLite session's database.""" - monkeypatch.setattr(workflow_response_converter, "db", SimpleNamespace(engine=sqlite_session.get_bind())) - return sqlite_session diff --git a/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py b/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py index 1a6dacff8e875f..ea1e852286e0f3 100644 --- a/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py @@ -8,21 +8,29 @@ from pytest_mock import MockerFixture from sqlalchemy.orm import Session -import core.app.apps.completion.app_generator as module -from core.app.apps.completion.app_generator import CompletionAppGenerator +import services.app.generation.adapters.completion as module from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom +from extensions.application_services.retrieval import build_dataset_retrieval from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from models import Account from models.enums import ConversationFromSource from models.model import App, AppMode, AppModelConfig, Conversation, Message +from services.app.generation.adapters.completion import CompletionAppGenerator +from services.errors.app import MoreLikeThisDisabledError +from services.errors.app_model_config import AppModelConfigBrokenError +from services.errors.message import MessageNotExistsError -TABLES = (AppModelConfig, Conversation, Message) +TABLES = (App, AppModelConfig, Conversation, Message) @pytest.fixture -def generator(mocker: MockerFixture): - gen = CompletionAppGenerator() +def generator(app_records, mocker: MockerFixture, *, annotation_replies, sqlite_session_factory): + gen = CompletionAppGenerator( + retrieval=build_dataset_retrieval(sqlite_session_factory), + annotations=annotation_replies, + records=app_records, + ) mocker.patch.object(module, "copy_current_request_context", side_effect=lambda fn: fn) @@ -40,14 +48,10 @@ def generator(mocker: MockerFixture): ) gen.generate_entity = generate_entity + mocker.patch.object(gen._records, "annotation_config", return_value={"enabled": False}) return gen -@pytest.fixture(autouse=True) -def _bind_db_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(module.db, "session", sqlite_session) - - def _build_app_model() -> App: return App( id="app1", @@ -78,6 +82,9 @@ def _persist_message( app_id: str = "app1", app_model_config: AppModelConfig | None = None, ) -> Message: + if session.get(App, app_id) is None: + session.add(_build_app_model()) + session.flush() conversation = Conversation( app_id=app_id, app_model_config_id=app_model_config.id if app_model_config else None, @@ -126,7 +133,7 @@ def _persist_message( app_mode=AppMode.COMPLETION, ) session.add(message) - session.flush() + session.commit() return message @@ -160,8 +167,8 @@ def test_generate_success_no_file_config(self, generator, mocker: MockerFixture, to_dict = mocker.patch.object(AppModelConfig, "to_dict", return_value={"model": {"provider": "x"}}) annotation_reply = {"enabled": False} load_annotation_reply_config = mocker.patch.object( - module, - "load_annotation_reply_config", + generator._records, + "annotation_config", return_value=annotation_reply, ) mocker.patch.object(module.FileUploadConfigManager, "convert", return_value=None) @@ -215,7 +222,7 @@ def test_generate_success_no_file_config(self, generator, mocker: MockerFixture, assert result == "converted" assert generator.generate_entity.call_args.kwargs["extras"]["trace_session_id"] == "session-1" module.file_factory.build_from_mappings.assert_not_called() - load_annotation_reply_config.assert_called_once_with(sqlite_session, "app1") + load_annotation_reply_config.assert_called_once_with(tenant_id="tenant", app_id="app1") to_dict.assert_called_once_with(annotation_reply=annotation_reply) assert get_app_config.call_args.kwargs["annotation_reply"] is annotation_reply @@ -329,6 +336,146 @@ def test_generate_override_model_config_debugger(self, generator, mocker: Mocker assert get_app_config.call_args.kwargs["override_config_dict"] == override_config + def test_generate_more_like_this_message_not_found(self, generator, sqlite_session: Session): + _persist_message(sqlite_session, app_id="other-app") + + with pytest.raises(MessageNotExistsError): + generator.generate_more_like_this( + session=sqlite_session, + app_model=_build_app_model(), + message_id="msg", + user=_build_user(), + invoke_from=InvokeFrom.WEB_APP, + ) + + def test_generate_more_like_this_disabled(self, generator, sqlite_session: Session): + app_model = _build_app_model() + current_config = AppModelConfig(app_id=app_model.id, more_like_this='{"enabled": false}') + sqlite_session.add(current_config) + sqlite_session.flush() + app_model.app_model_config_id = current_config.id + _persist_message(sqlite_session) + + with pytest.raises(MoreLikeThisDisabledError): + generator.generate_more_like_this( + session=sqlite_session, + app_model=app_model, + message_id="msg", + user=_build_user(), + invoke_from=InvokeFrom.WEB_APP, + ) + + def test_generate_more_like_this_app_model_config_missing(self, generator, sqlite_session: Session): + app_model = _build_app_model() + app_model.app_model_config_id = None + + _persist_message(sqlite_session) + + with pytest.raises(MoreLikeThisDisabledError): + generator.generate_more_like_this( + session=sqlite_session, + app_model=app_model, + message_id="msg", + user=_build_user(), + invoke_from=InvokeFrom.WEB_APP, + ) + + def test_generate_more_like_this_message_config_none(self, generator, sqlite_session: Session): + app_model = _build_app_model() + current_config = AppModelConfig(app_id=app_model.id, more_like_this='{"enabled": true}') + sqlite_session.add(current_config) + sqlite_session.flush() + app_model.app_model_config_id = current_config.id + _persist_message(sqlite_session) + + with pytest.raises(AppModelConfigBrokenError): + generator.generate_more_like_this( + session=sqlite_session, + app_model=app_model, + message_id="msg", + user=_build_user(), + invoke_from=InvokeFrom.WEB_APP, + ) + + def test_generate_more_like_this_success(self, generator, mocker: MockerFixture, sqlite_session: Session): + app_model = _build_app_model() + app_model_config = AppModelConfig(app_id=app_model.id, more_like_this='{"enabled": true}') + sqlite_session.add(app_model_config) + sqlite_session.flush() + app_model.app_model_config_id = app_model_config.id + _persist_message(sqlite_session, app_model_config=app_model_config) + to_dict = mocker.patch.object( + AppModelConfig, + "to_dict", + return_value={ + "model": {"completion_params": {"temperature": 0.1}}, + "file_upload": {"enabled": True}, + }, + ) + annotation_reply = {"enabled": False} + load_annotation_reply_config = mocker.patch.object( + generator._records, + "annotation_config", + return_value=annotation_reply, + ) + mocker.patch.object(module.FileUploadConfigManager, "convert", return_value=None) + build_from_mappings = mocker.patch.object(module.file_factory, "build_from_mappings") + + app_config = MagicMock(variables=["v"], to_dict=MagicMock(return_value={})) + get_app_config = mocker.patch.object( + module.CompletionAppConfigManager, + "get_app_config", + return_value=app_config, + ) + mocker.patch.object(module.ModelConfigConverter, "convert", return_value=MagicMock()) + + mocker.patch.object( + generator, + "_init_generate_records", + return_value=( + Conversation( + id="conv", + app_id="app1", + mode=AppMode.COMPLETION, + name="Conversation", + _inputs={}, + from_source=ConversationFromSource.CONSOLE, + ), + Message( + id="msg", + app_id="app1", + conversation_id="conv", + _inputs={}, + query="q", + message={}, + message_unit_price=Decimal(0), + answer="", + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ), + ), + ) + mocker.patch.object(generator, "_handle_response", return_value="response") + mocker.patch.object(module.CompletionAppGenerateResponseConverter, "convert", return_value="converted") + + result = generator.generate_more_like_this( + session=sqlite_session, + app_model=app_model, + message_id="msg", + user=_build_user(), + invoke_from=InvokeFrom.WEB_APP, + stream=True, + ) + + assert result == "converted" + load_annotation_reply_config.assert_called_once_with(tenant_id="tenant", app_id="app1") + to_dict.assert_called_once_with(annotation_reply=annotation_reply) + assert generator.generate_entity.call_args.kwargs["inputs"] == {"a": 1} + override_dict = get_app_config.call_args.kwargs["override_config_dict"] + assert override_dict["model"]["completion_params"]["temperature"] == 0.9 + build_from_mappings.assert_not_called() + @pytest.mark.parametrize( ("error", "should_publish"), [ @@ -357,7 +504,7 @@ def test_generate_worker_error_handling( flask_app.app_context.return_value = contextlib.nullcontext() message = _persist_message(sqlite_session) - mocker.patch.object(generator, "_get_message", return_value=message) + mocker.patch.object(generator._records, "load", return_value=(_build_app_model(), None, message)) runner_instance = MagicMock() runner_instance.run.side_effect = error @@ -369,6 +516,7 @@ def test_generate_worker_error_handling( application_generate_entity=MagicMock(), queue_manager=queue_manager, message_id="msg", + conversation_id="conv", ) assert queue_manager.publish_error.called is should_publish diff --git a/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py b/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py index 6a463c05e23f40..cdf81275c73162 100644 --- a/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py @@ -1,13 +1,11 @@ +import json from collections.abc import Callable, Generator, Iterator from dataclasses import dataclass, field from decimal import Decimal import pytest -from flask import Flask from redis import Redis -from sqlalchemy import event -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.app.app_config.entities import ( AppAdditionalFeatures, @@ -19,7 +17,6 @@ PromptTemplateEntity, ) from core.app.apps.completion.app_config_manager import CompletionAppConfig -from core.app.apps.completion.app_runner import CompletionAppRunner from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import CompletionAppGenerateEntity, InvokeFrom from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy @@ -31,7 +28,7 @@ from core.plugin.impl.model_runtime import PluginModelRuntime from core.plugin.impl.model_runtime_factory import create_model_type_instance, create_plugin_model_runtime from core.provider_manager import ProviderManager -from extensions.ext_database import db +from extensions.application_services.retrieval import build_dataset_retrieval from extensions.ext_redis import RedisClientWrapper from graphon.file import FileUploadConfig from graphon.file.models import ImageConfig @@ -42,10 +39,13 @@ UserPromptMessage, ) from graphon.model_runtime.entities.model_entities import ModelType -from models.dataset import Dataset +from models.dataset import Dataset, ExternalKnowledgeApis, ExternalKnowledgeBindings from models.enums import ConversationFromSource from models.enums import InvokeFrom as ModelInvokeFrom from models.model import App, AppMode, Conversation, IconType, Message +from repositories.app.generation_repository import AppGenerationRepository +from services.app.generation.adapters.completion_runner import CompletionAppRunner +from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalKnowledgeApiSetting from services.knowledge.external.service import ExternalDatasetService from tests.helpers.unit.core.model_fixtures import make_model_config @@ -152,21 +152,13 @@ def dispatch(_client: PluginModelClient, **kwargs: object) -> Generator[LLMResul @pytest.fixture -def retrieval_app(sqlite_engine: Engine) -> Iterator[Flask]: - """Bind the real Flask database extension for independent retrieval audit writes.""" - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = sqlite_engine.url - db.init_app(app) - with app.app_context(): - try: - yield app - finally: - db.engine.dispose() - - -@pytest.fixture -def runner() -> CompletionAppRunner: - return CompletionAppRunner() +def runner( + *, app_records: AppGenerationRepository, sqlite_session_factory: sessionmaker[Session] +) -> CompletionAppRunner: + return CompletionAppRunner( + retrieval=build_dataset_retrieval(sqlite_session_factory), + records=app_records, + ) def _build_app_config( @@ -263,19 +255,10 @@ def _persist_records(session: Session) -> tuple[App, Message]: class TestCompletionAppRunner: - def test_run_app_not_found( - self, runner: CompletionAppRunner, sqlite_session: Session, queue_manager: MessageBasedAppQueueManager - ) -> None: - app_config = _build_app_config() - app_generate_entity = _build_generate_entity(app_config) - - with pytest.raises(ValueError, match="App not found"): - runner.run(app_generate_entity, queue_manager, _message(), sqlite_session) - def test_run_moderation_error_outputs_direct( self, runner: CompletionAppRunner, sqlite_session: Session, queue_manager: MessageBasedAppQueueManager ) -> None: - _, message = _persist_records(sqlite_session) + app_record, message = _persist_records(sqlite_session) app_config = _build_app_config() app_generate_entity = _build_generate_entity(app_config) @@ -287,7 +270,7 @@ def test_run_moderation_error_outputs_direct( runner.direct_output = direct_output # type: ignore[method-assign] runner._handle_invoke_result = handle_invoke_result # type: ignore[method-assign] - runner.run(app_generate_entity, queue_manager, message, sqlite_session) + runner.run(app_generate_entity, queue_manager, message, app_record=app_record) assert len(direct_output.calls) == 1 assert handle_invoke_result.calls == [] @@ -295,7 +278,7 @@ def test_run_moderation_error_outputs_direct( def test_run_hosting_moderation_stops( self, runner: CompletionAppRunner, sqlite_session: Session, queue_manager: MessageBasedAppQueueManager ) -> None: - _, message = _persist_records(sqlite_session) + app_record, message = _persist_records(sqlite_session) app_config = _build_app_config() app_generate_entity = _build_generate_entity(app_config) @@ -308,7 +291,7 @@ def test_run_hosting_moderation_stops( handle_invoke_result = CallRecorder() runner._handle_invoke_result = handle_invoke_result # type: ignore[method-assign] - runner.run(app_generate_entity, queue_manager, message, sqlite_session) + runner.run(app_generate_entity, queue_manager, message, app_record=app_record) assert handle_invoke_result.calls == [] @@ -319,13 +302,35 @@ def test_run_dataset_and_external_tools_flow( sqlite_session: Session, queue_manager: MessageBasedAppQueueManager, model_runtime: tuple[ProviderModelBundle, list[RecordedCall], list[str]], - retrieval_app: Flask, ) -> None: - _, message = _persist_records(sqlite_session) + app_record, message = _persist_records(sqlite_session) dataset = Dataset( - id="ds", tenant_id=TENANT_ID, name="External knowledge", provider="external", created_by="user" + id="ds", + tenant_id=TENANT_ID, + name="External knowledge", + provider="external", + retrieval_model={"top_k": 3}, + created_by="user", + ) + external_api = ExternalKnowledgeApis( + name="External API", + description="", + tenant_id=TENANT_ID, + settings=json.dumps({"endpoint": "https://knowledge.example", "api_key": "secret"}), + created_by="user", + updated_by="user", + ) + sqlite_session.add_all([dataset, external_api]) + sqlite_session.flush() + sqlite_session.add( + ExternalKnowledgeBindings( + tenant_id=TENANT_ID, + external_knowledge_api_id=external_api.id, + dataset_id=dataset.id, + external_knowledge_id="external-dataset", + created_by="user", + ) ) - sqlite_session.add(dataset) sqlite_session.commit() retrieve_config = DatasetRetrieveConfigEntity( @@ -359,78 +364,30 @@ def test_run_dataset_and_external_tools_flow( handle_invoke_result = CallRecorder() runner._handle_invoke_result = handle_invoke_result # type: ignore[method-assign] - retrieval_calls: list[dict[str, object]] = [] + retrieval_calls: list[ExternalKnowledgeApiSetting] = [] - def retrieve_external(**kwargs: object) -> list[dict[str, object]]: - retrieval_calls.append(kwargs) + def retrieve_external(request: ExternalKnowledgeApiSetting) -> list[dict[str, object]]: + retrieval_calls.append(request) return [{"content": "ctx", "metadata": {}, "score": 1.0, "title": "Document"}] - monkeypatch.setattr(ExternalDatasetService, "fetch_external_knowledge_retrieval", retrieve_external) + monkeypatch.setattr(ExternalDatasetService, "execute_external_knowledge_retrieval", retrieve_external) - runner.run(app_generate_entity, queue_manager, message, sqlite_session) + runner.run(app_generate_entity, queue_manager, message, app_record=app_record) assert len(retrieval_calls) == 1 - assert retrieval_calls[0]["query"] == "query_from_input" - assert retrieval_calls[0]["tenant_id"] == TENANT_ID - assert retrieval_calls[0]["dataset_id"] == dataset.id + request = retrieval_calls[0] + assert request.url == "https://knowledge.example/retrieval" + assert request.headers == {"Content-Type": "application/json", "Authorization": "Bearer secret"} + assert request.params == { + "retrieval_setting": {"top_k": 3, "score_threshold": 0.0}, + "query": "query_from_input", + "knowledge_id": "external-dataset", + "metadata_condition": None, + } assert organize_prompt_messages.calls[-1].kwargs["context"] == "ctx" assert organize_prompt_messages.calls[-1].kwargs["context_files"] == [] assert len(handle_invoke_result.calls) == 1 - def test_run_closes_explicit_session_before_stream_consumption( - self, - runner: CompletionAppRunner, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - queue_manager: MessageBasedAppQueueManager, - model_runtime: tuple[ProviderModelBundle, list[RecordedCall], list[str]], - ) -> None: - _, message = _persist_records(sqlite_session) - app_config = _build_app_config() - app_generate_entity = _build_generate_entity(app_config) - - _, _, events = model_runtime - session = sqlite_session - original_close = session.close - - def record_commit(_session: Session) -> None: - events.append("commit") - - def record_close() -> None: - events.append("close") - original_close() - - runner.organize_prompt_messages = CallRecorder(result=([], None)) # type: ignore[method-assign] - runner.moderation_for_inputs = CallRecorder( # type: ignore[method-assign] - result=(None, app_generate_entity.inputs, "query") - ) - runner.check_hosting_moderation = CallRecorder(result=False) # type: ignore[method-assign] - runner.recalc_llm_max_tokens = CallRecorder() # type: ignore[method-assign] - handle_invoke_result = CallRecorder(implementation=lambda invoke_result, **kwargs: list(invoke_result)) - runner._handle_invoke_result = handle_invoke_result # type: ignore[method-assign] - - monkeypatch.setattr(session, "close", record_close) - - event.listen(session, "after_commit", record_commit) - try: - runner.run(app_generate_entity, queue_manager, message, session) - finally: - event.remove(session, "after_commit", record_commit) - - assert events == ["commit", "close", "invoke", "first-chunk"] - assert len(handle_invoke_result.calls) == 1 - handle_call = handle_invoke_result.calls[0] - assert handle_call.args == () - assert handle_call.kwargs["queue_manager"] is queue_manager - assert handle_call.kwargs | {"invoke_result": None} == { - "invoke_result": None, - "queue_manager": queue_manager, - "stream": True, - "message_id": "msg", - "user_id": "user", - "tenant_id": TENANT_ID, - } - @pytest.mark.parametrize("stream", [False, True]) def test_run_invokes_model_resolved_by_model_manager( self, @@ -440,7 +397,7 @@ def test_run_invokes_model_resolved_by_model_manager( queue_manager: MessageBasedAppQueueManager, model_runtime: tuple[ProviderModelBundle, list[RecordedCall], list[str]], ) -> None: - _, message = _persist_records(sqlite_session) + app_record, message = _persist_records(sqlite_session) app_config = _build_app_config() app_generate_entity = _build_generate_entity(app_config) app_generate_entity.stream = stream @@ -460,7 +417,7 @@ def handle_result(invoke_result: LLMResult | Generator[LLMResultChunk], **_kwarg runner._handle_invoke_result = handle_result # type: ignore[method-assign] - runner.run(app_generate_entity, queue_manager, message, sqlite_session) + runner.run(app_generate_entity, queue_manager, message, app_record=app_record) bundle, calls, _ = model_runtime assert bundle.configuration.tenant_id == TENANT_ID @@ -512,7 +469,7 @@ def handle_result(invoke_result: LLMResult | Generator[LLMResultChunk], **_kwarg def test_run_uses_low_image_detail_default( self, runner: CompletionAppRunner, sqlite_session: Session, queue_manager: MessageBasedAppQueueManager ) -> None: - _, message = _persist_records(sqlite_session) + app_record, message = _persist_records(sqlite_session) app_config = _build_app_config() app_generate_entity = _build_generate_entity(app_config, file_upload_config=None) @@ -524,6 +481,6 @@ def test_run_uses_low_image_detail_default( ) runner.check_hosting_moderation = CallRecorder(result=True) # type: ignore[method-assign] - runner.run(app_generate_entity, queue_manager, message, sqlite_session) + runner.run(app_generate_entity, queue_manager, message, app_record=app_record) assert organize_prompt_messages.calls[-1].kwargs["image_detail_config"] == ImagePromptMessageContent.DETAIL.LOW diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_generate_response_converter.py b/api/tests/unit_tests/core/app/apps/pipeline/test_generate_response_converter.py index 62e9125a5f4650..49cc7cc72ef1fe 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_generate_response_converter.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_generate_response_converter.py @@ -1,6 +1,5 @@ from collections.abc import Generator -from core.app.apps.pipeline.generate_response_converter import WorkflowAppGenerateResponseConverter from core.app.entities.task_entities import ( AppStreamResponse, ErrorStreamResponse, @@ -11,6 +10,9 @@ WorkflowAppStreamResponse, ) from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from services.workflow.execution.adapters.pipeline.generate_response_converter import ( + WorkflowAppGenerateResponseConverter, +) def test_convert_blocking_full_and_simple_response() -> None: diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py index 64347ec775ae1a..8369caa37ae2c5 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py @@ -3,11 +3,12 @@ from pytest_mock import MockerFixture -import core.app.apps.pipeline.pipeline_config_manager as module -from core.app.apps.pipeline.pipeline_config_manager import PipelineConfigManager +import services.workflow.execution.adapters.pipeline.pipeline_config_manager as module from models.dataset import Pipeline from models.model import AppMode from models.workflow import Workflow, WorkflowType +from repositories.workflow.definition_repository import workflow_snapshot +from services.workflow.execution.adapters.pipeline.pipeline_config_manager import PipelineConfigManager def test_get_pipeline_config(mocker: MockerFixture) -> None: @@ -34,7 +35,9 @@ def test_get_pipeline_config(mocker: MockerFixture) -> None: ) mocker.patch.object(module, "PipelineConfig", side_effect=lambda **kwargs: SimpleNamespace(**kwargs)) - result = PipelineConfigManager.get_pipeline_config(pipeline=pipeline, workflow=workflow, start_node_id="start") + result = PipelineConfigManager.get_pipeline_config( + pipeline=pipeline, workflow=workflow_snapshot(workflow), start_node_id="start" + ) assert result.tenant_id == "tenant" assert result.app_id == "pipe1" diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_queue_manager.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_queue_manager.py index 0adcfc9e6e9b69..c7a38ca0542c2a 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_queue_manager.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_queue_manager.py @@ -3,7 +3,6 @@ from core.app.apps.base_app_queue_manager import PublishFrom from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.pipeline.pipeline_queue_manager import PipelineQueueManager from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import ( QueueErrorEvent, @@ -15,6 +14,7 @@ QueueWorkflowSucceededEvent, ) from graphon.model_runtime.entities.llm_entities import LLMResult +from services.workflow.execution.adapters.pipeline.pipeline_queue_manager import PipelineQueueManager def test_publish_sets_stop_listen_and_raises_on_stopped(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py index 4c7e8a1fbad56b..aad24eefb8f4c5 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py @@ -1,7 +1,7 @@ """Unit tests for PipelineRunner behavior. This module validates core control-flow outcomes for -``core.app.apps.pipeline.pipeline_runner``: app/workflow lookup, graph +``services.workflow.execution.adapters.pipeline.pipeline_runner``: app/workflow lookup, graph initialization guards, invoke-source to user-source resolution, and failed-run event handling. Invariants asserted here include strict graph-config validation, correct ``InvokeFrom`` to ``UserFrom`` mapping, and publishing @@ -11,21 +11,16 @@ ``UserFrom``, and real ORM, graph, and runtime dependencies used by the runner. """ -import json from collections.abc import Generator import pytest +from pytest_mock import MockerFixture from sqlalchemy import event from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker -import core.app.apps.pipeline.pipeline_runner as module -from core.app.apps.pipeline.pipeline_config_manager import PipelineConfig -from core.app.apps.pipeline.pipeline_queue_manager import PipelineQueueManager -from core.app.apps.pipeline.pipeline_runner import PipelineRunner from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity, UserFrom from core.app.entities.queue_entities import AppQueueEvent, QueueWorkflowFailedEvent -from core.workflow.workflow_entry import WorkflowEntry from graphon.engine_events import EngineEvent, GraphRunFailedEvent from graphon.graph import Graph from graphon.runtime import RuntimeState, VariablePool @@ -35,6 +30,14 @@ from models.model import AppMode, EndUser from models.workflow import Workflow, WorkflowType from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from repositories.workflow.definition_repository import workflow_snapshot +from services.errors.app import WorkflowNotFoundError +from services.workflow.contracts import WorkflowSnapshot +from services.workflow.execution.adapters.pipeline.pipeline_config_manager import PipelineConfig +from services.workflow.execution.adapters.pipeline.pipeline_queue_manager import PipelineQueueManager +from services.workflow.execution.adapters.pipeline.pipeline_runner import PipelineRunner +from services.workflow.execution.adapters.workflow_entry import WorkflowEntry +from services.workflow.execution.ports import WorkflowRuntime from tests.helpers.unit.model_factories import make_dataset, make_document, make_end_user, make_workflow @@ -55,7 +58,9 @@ def _dataset(*, tenant_id: str = "tenant", dataset_id: str = "ds", pipeline_id: ) -def _workflow(*, tenant_id: str = "tenant", pipeline_id: str = "pipe", graph: dict | None = None) -> Workflow: +def _workflow( + *, tenant_id: str = "tenant", pipeline_id: str = "pipe", graph: dict[str, object] | None = None +) -> Workflow: return make_workflow( tenant_id=tenant_id, app_id=pipeline_id, @@ -132,7 +137,11 @@ def isolate_queue_redis(monkeypatch: pytest.MonkeyPatch) -> None: def _build_runner( - app_generate_entity: RagPipelineGenerateEntity, sqlite_engine: Engine, workflow: Workflow + app_generate_entity: RagPipelineGenerateEntity, + sqlite_engine: Engine, + workflow: Workflow, + *, + workflow_runtime: WorkflowRuntime, ) -> PipelineRunner: return PipelineRunner( application_generate_entity=app_generate_entity, @@ -143,56 +152,48 @@ def _build_runner( app_mode=AppMode.RAG_PIPELINE, ), variable_loader=DUMMY_VARIABLE_LOADER, - workflow=workflow, + workflow=workflow_snapshot(workflow), system_user_id="sys", documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), + runtime=workflow_runtime, ) @pytest.fixture -def runner(sqlite_engine: Engine): - return _build_runner(_build_app_generate_entity(), sqlite_engine, _workflow()) - +def runner(sqlite_engine: Engine, *, workflow_runtime: WorkflowRuntime) -> PipelineRunner: + return _build_runner(_build_app_generate_entity(), sqlite_engine, _workflow(), workflow_runtime=workflow_runtime) -def test_get_app_id(runner): - assert runner._get_app_id() == "pipe" - -def test_get_workflow_returns_workflow(runner, sqlite_session: Session): +def test_get_workflow_returns_workflow(runner: PipelineRunner, sqlite_session: Session) -> None: pipeline, _, workflow = _persist_scope(sqlite_session) - result = runner.get_workflow(session=sqlite_session, pipeline=pipeline, workflow_id="wf") + result = runner._runtime.contexts.workflow(tenant_id=pipeline.tenant_id, app_id=pipeline.id, workflow_id="wf") - assert result == workflow + assert result.id == workflow.id -def test_init_rag_pipeline_graph_invalid_config(mocker, runner): - workflow = _workflow(graph={}) +def test_init_rag_pipeline_graph_invalid_config(runner: PipelineRunner) -> None: + state = RuntimeState(workflow_id="wf", variable_pool=VariablePool(), start_at=0) - with pytest.raises(ValueError, match="nodes or edges not found in workflow graph"): - runner._init_rag_pipeline_graph( - workflow=workflow, - graph_runtime_state=RuntimeState(workflow_id=workflow.id, variable_pool=VariablePool(), start_at=0), - ) + with pytest.raises(ValueError, match="validation error"): + runner._init_rag_pipeline_graph(workflow=workflow_snapshot(_workflow(graph={})), graph_runtime_state=state) - workflow.graph = json.dumps({"nodes": "bad", "edges": []}) - with pytest.raises(ValueError, match="nodes in workflow graph must be a list"): + with pytest.raises(ValueError, match="validation error"): runner._init_rag_pipeline_graph( - workflow=workflow, - graph_runtime_state=RuntimeState(workflow_id=workflow.id, variable_pool=VariablePool(), start_at=0), + workflow=workflow_snapshot(_workflow(graph={"nodes": "bad", "edges": []})), + graph_runtime_state=state, ) - workflow.graph = json.dumps({"nodes": [], "edges": "bad"}) - with pytest.raises(ValueError, match="edges in workflow graph must be a list"): + with pytest.raises(ValueError, match="validation error"): runner._init_rag_pipeline_graph( - workflow=workflow, - graph_runtime_state=RuntimeState(workflow_id=workflow.id, variable_pool=VariablePool(), start_at=0), + workflow=workflow_snapshot(_workflow(graph={"nodes": [], "edges": "bad"})), + graph_runtime_state=state, ) -def test_init_rag_pipeline_graph_not_found(mocker, runner): - workflow = _workflow() - mocker.patch.object(module.Graph, "init", return_value=None) +def test_init_rag_pipeline_graph_not_found(mocker: MockerFixture, runner: PipelineRunner) -> None: + workflow = workflow_snapshot(_workflow()) + mocker.patch("services.workflow.execution.adapters.graph.Graph.init", return_value=None) with pytest.raises(ValueError, match="Unable to determine default root node ID from workflow graph"): runner._init_rag_pipeline_graph( @@ -201,7 +202,7 @@ def test_init_rag_pipeline_graph_not_found(mocker, runner): ) -def test_init_rag_pipeline_graph_honors_canonical_root_ownership(runner): +def test_init_rag_pipeline_graph_honors_canonical_root_ownership(runner: PipelineRunner) -> None: workflow = _workflow( graph={ "nodes": [ @@ -217,13 +218,13 @@ def test_init_rag_pipeline_graph_honors_canonical_root_ownership(runner): ) runtime_state = RuntimeState(workflow_id=workflow.id, variable_pool=VariablePool(), start_at=0.0) - graph = runner._init_rag_pipeline_graph(workflow=workflow, graph_runtime_state=runtime_state) + graph = runner._init_rag_pipeline_graph(workflow=workflow_snapshot(workflow), graph_runtime_state=runtime_state) assert graph.root_node.id == "start" assert set(graph.nodes) == {"start", "end"} -def test_update_document_status_on_failure(runner, sqlite_session: Session): +def test_update_document_status_on_failure(runner: PipelineRunner, sqlite_session: Session) -> None: document = _document() _, dataset, _ = _persist_scope(sqlite_session, documents=(document,)) @@ -243,7 +244,7 @@ def test_update_document_status_on_failure(runner, sqlite_session: Session): assert updated.error == "boom" -def test_update_document_status_skips_when_document_not_found(runner, sqlite_session: Session): +def test_update_document_status_skips_when_document_not_found(runner: PipelineRunner, sqlite_session: Session) -> None: _, dataset, _ = _persist_scope(sqlite_session) runner._update_document_status( @@ -256,10 +257,10 @@ def test_update_document_status_skips_when_document_not_found(runner, sqlite_ses assert sqlite_session.get(Document, "missing") is None -def test_update_document_status_skips_without_document_ref(runner, sqlite_engine: Engine): +def test_update_document_status_skips_without_document_ref(runner: PipelineRunner, sqlite_engine: Engine) -> None: checkouts = 0 - def record_checkout(*_args) -> None: + def record_checkout(*_args: object) -> None: nonlocal checkouts checkouts += 1 @@ -277,85 +278,75 @@ def record_checkout(*_args) -> None: assert checkouts == 0 -def test_run_pipeline_not_found(sqlite_engine: Engine): +def test_run_pipeline_not_found(sqlite_engine: Engine, *, workflow_runtime: WorkflowRuntime) -> None: app_generate_entity = _build_app_generate_entity() app_generate_entity.invoke_from = InvokeFrom.WEB_APP app_generate_entity.single_iteration_run = None app_generate_entity.single_loop_run = None - runner = _build_runner(app_generate_entity, sqlite_engine, _workflow()) + runner = _build_runner(app_generate_entity, sqlite_engine, _workflow(), workflow_runtime=workflow_runtime) - with pytest.raises(ValueError, match="Pipeline not found"): + with pytest.raises(WorkflowNotFoundError): runner.prepare() -def test_run_pipeline_from_other_tenant_is_not_found(runner: PipelineRunner, sqlite_session: Session): +def test_run_pipeline_from_other_tenant_is_not_found(runner: PipelineRunner, sqlite_session: Session) -> None: pipeline = _pipeline(tenant_id="other-tenant") sqlite_session.add(pipeline) sqlite_session.commit() - with pytest.raises(ValueError, match="Pipeline not found"): + with pytest.raises(WorkflowNotFoundError, match="Pipeline not found"): runner.prepare() -def _unexpected_workflow_lookup(session: Session, pipeline: Pipeline, workflow_id: str) -> Workflow | None: - pytest.fail("Workflow lookup must follow successful dataset and document ownership checks") - - @pytest.mark.parametrize( - "dataset", + ("dataset", "error_type"), [ - pytest.param(None, id="missing"), - pytest.param(_dataset(tenant_id="other-tenant"), id="other-tenant"), - pytest.param(_dataset(dataset_id="other-dataset"), id="other-dataset"), + pytest.param(None, WorkflowNotFoundError, id="missing"), + pytest.param(_dataset(tenant_id="other-tenant"), WorkflowNotFoundError, id="other-tenant"), + pytest.param(_dataset(dataset_id="other-dataset"), ValueError, id="other-dataset"), ], ) def test_run_rejects_unowned_pipeline_dataset( - monkeypatch: pytest.MonkeyPatch, runner: PipelineRunner, dataset: Dataset | None, + error_type: type[Exception], sqlite_session: Session, -): +) -> None: pipeline = _pipeline() sqlite_session.add(pipeline) if dataset is not None: sqlite_session.add(dataset) sqlite_session.commit() - monkeypatch.setattr(runner, "get_workflow", _unexpected_workflow_lookup) - - with pytest.raises(ValueError, match="Pipeline dataset not found"): + with pytest.raises(error_type, match="Pipeline dataset not found"): runner.prepare() def test_run_rejects_document_outside_pipeline_dataset_after_async_boundary( - monkeypatch: pytest.MonkeyPatch, runner: PipelineRunner, sqlite_session: Session, -): +) -> None: runner.application_generate_entity.document_id = "foreign-doc" runner.application_generate_entity.original_document_id = "foreign-doc" _persist_scope(sqlite_session) - monkeypatch.setattr(runner, "get_workflow", _unexpected_workflow_lookup) - with pytest.raises(ValueError, match="Pipeline document not found"): runner.prepare() def test_run_rejects_original_document_outside_pipeline_dataset_after_async_boundary( - monkeypatch: pytest.MonkeyPatch, runner: PipelineRunner, sqlite_session: Session, -): +) -> None: runner.application_generate_entity.document_id = "doc" runner.application_generate_entity.original_document_id = "foreign-doc" _persist_scope(sqlite_session, documents=(_document(),)) - monkeypatch.setattr(runner, "get_workflow", _unexpected_workflow_lookup) - - with pytest.raises(ValueError, match="Pipeline original document not found"): + with pytest.raises(ValueError, match="Pipeline document not found"): runner.prepare() -def test_run_workflow_not_initialized(sqlite_session: Session, sqlite_engine: Engine): +def test_run_rejects_invalid_workflow_snapshot( + sqlite_session: Session, sqlite_engine: Engine, *, workflow_runtime: WorkflowRuntime +) -> None: app_generate_entity = _build_app_generate_entity() pipeline = _pipeline() @@ -364,8 +355,8 @@ def test_run_workflow_not_initialized(sqlite_session: Session, sqlite_engine: En sqlite_session.add_all([pipeline, dataset, document]) sqlite_session.commit() - runner = _build_runner(app_generate_entity, sqlite_engine, _workflow()) - with pytest.raises(ValueError, match="Workflow not initialized"): + runner = _build_runner(app_generate_entity, sqlite_engine, _workflow(), workflow_runtime=workflow_runtime) + with pytest.raises(ValueError, match="Graph must have at least one node"): runner.prepare() @@ -375,14 +366,27 @@ def _start_workflow() -> Workflow: ) -def test_run_single_iteration_path(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine): +def test_run_single_iteration_path( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: Engine, + *, + workflow_runtime: WorkflowRuntime, +) -> None: app_generate_entity = _build_app_generate_entity() app_generate_entity.single_iteration_run = RagPipelineGenerateEntity.SingleIterationRunEntity( node_id="start", inputs={} ) _, dataset, workflow = _persist_scope(sqlite_session, workflow=_start_workflow(), documents=(_document(),)) - runner = _build_runner(app_generate_entity, sqlite_engine, workflow) - prepared = [] + runner = _build_runner(app_generate_entity, sqlite_engine, workflow, workflow_runtime=workflow_runtime) + prepared: list[ + tuple[ + str, + RagPipelineGenerateEntity.SingleIterationRunEntity | None, + RagPipelineGenerateEntity.SingleLoopRunEntity | None, + str, + ] + ] = [] published: list[AppQueueEvent] = [] failure = GraphRunFailedEvent(error="iteration failed") variable_pool = VariablePool() @@ -390,7 +394,7 @@ def test_run_single_iteration_path(monkeypatch: pytest.MonkeyPatch, sqlite_sessi def prepare_single_node_execution( *, - workflow: Workflow, + workflow: WorkflowSnapshot, single_iteration_run: RagPipelineGenerateEntity.SingleIterationRunEntity | None, single_loop_run: RagPipelineGenerateEntity.SingleLoopRunEntity | None, user_id: str, @@ -403,8 +407,8 @@ def run_workflow(entry: WorkflowEntry) -> Generator[EngineEvent]: assert entry.graph_engine.runtime_state is state yield failure - monkeypatch.setattr(runner, "_prepare_single_node_execution", prepare_single_node_execution) - monkeypatch.setattr(runner, "_publish_event", published.append) + monkeypatch.setattr(runner._graphs, "build_single_node", prepare_single_node_execution) + monkeypatch.setattr(runner._events, "_publish_event", published.append) monkeypatch.setattr(WorkflowEntry, "run", run_workflow) prepared_run = runner.prepare() @@ -425,7 +429,13 @@ def run_workflow(entry: WorkflowEntry) -> Generator[EngineEvent]: assert published[0].error == failure.error -def test_run_normal_path_builds_graph(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine): +def test_run_normal_path_builds_graph( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_engine: Engine, + *, + workflow_runtime: WorkflowRuntime, +) -> None: app_generate_entity = _build_app_generate_entity() events: list[str] = [] workflow = _start_workflow() @@ -433,7 +443,7 @@ def test_run_normal_path_builds_graph(monkeypatch: pytest.MonkeyPatch, sqlite_se {"variable": "input1", "belong_to_node_id": "start", "type": "text-input", "label": "Input"} ] _persist_scope(sqlite_session, workflow=workflow, end_user=_end_user(), documents=(_document(),)) - runner = _build_runner(app_generate_entity, sqlite_engine, workflow) + runner = _build_runner(app_generate_entity, sqlite_engine, workflow, workflow_runtime=workflow_runtime) entries: list[WorkflowEntry] = [] def run_workflow(entry: WorkflowEntry) -> Generator[EngineEvent]: @@ -443,7 +453,7 @@ def run_workflow(entry: WorkflowEntry) -> Generator[EngineEvent]: monkeypatch.setattr(WorkflowEntry, "run", run_workflow) - def record_checkin(*_args) -> None: + def record_checkin(*_args: object) -> None: events.append("session_checkin") event.listen(sqlite_engine, "checkin", record_checkin) 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 c7735b7e834047..bdcd20eb24aa9e 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,13 +1,11 @@ import logging from copy import deepcopy from datetime import UTC, datetime -from unittest.mock import Mock +from unittest.mock import Mock, call, create_autospec import pytest -from sqlalchemy import inspect -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker -from core.app.apps.base_app_generator import BaseAppGenerator from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.app.file_access import FileAccessScope, bind_file_access_scope, get_current_file_access_scope from core.workflow.file_reference import resolve_file_record_id @@ -15,10 +13,16 @@ from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus from graphon.file import File from graphon.variables.input_entities import VariableEntity, VariableEntityType -from models import CreatorUserRole, UploadFile, Workflow, WorkflowRun, WorkflowRunTriggeredFrom +from models import CreatorUserRole, UploadFile, WorkflowRun, WorkflowRunTriggeredFrom from models.workflow import WorkflowType as ModelWorkflowType +from repositories.workflow.definition_repository import workflow_snapshot +from repositories.workflow.runtime_context_repository import WorkflowRuntimeContextRepository from services.app.generation.input_adapter import AppInputAdapter +from services.app.generation.response import convert_to_event_stream +from services.errors.app import WorkflowNotFoundError +from services.workflow.execution.generation_service import consume_stream, join_worker from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_workflow def _workflow_run(*, graph: str | None) -> WorkflowRun: @@ -39,20 +43,25 @@ def _workflow_run(*, graph: str | None) -> WorkflowRun: def test_restore_workflow_run_graph(sqlite_session: Session): - workflow = Workflow(graph='{"nodes": [{"id": "edited"}]}') + workflow = make_workflow( + workflow_id="workflow-1", + tenant_id="tenant-1", + app_id="app-1", + graph='{"nodes": [{"id": "edited"}]}', + ) workflow_run = _workflow_run(graph='{"nodes": [{"id": "paused"}]}') sqlite_session.add(workflow_run) sqlite_session.commit() - BaseAppGenerator._restore_workflow_run_graph( - session=sqlite_session, - workflow=workflow, + snapshot = workflow_snapshot(workflow) + restored = WorkflowRuntimeContextRepository(sessionmaker(bind=sqlite_session.get_bind())).restore_graph( + workflow=snapshot, workflow_run_id="run-id", ) assert sqlite_session.get(WorkflowRun, "run-id") is workflow_run - assert workflow.graph == '{"nodes": [{"id": "paused"}]}' - assert not inspect(workflow).attrs.graph.history.has_changes() + assert restored.graph == '{"nodes": [{"id": "paused"}]}' + assert snapshot.graph == '{"nodes": [{"id": "edited"}]}' @pytest.mark.parametrize( @@ -69,17 +78,19 @@ def test_restore_workflow_run_graph_requires_persisted_snapshot( sqlite_session.commit() with pytest.raises( - ValueError, match=("(?:Workflow run graph not found: run\\-id|Workflow run id is required when resuming)") + WorkflowNotFoundError, + match="(?:Workflow run graph not found: run\\-id|Workflow run id is required when resuming)", ): - BaseAppGenerator._restore_workflow_run_graph( - session=sqlite_session, - workflow=Workflow(graph="{}"), + WorkflowRuntimeContextRepository(sessionmaker(bind=sqlite_session.get_bind())).restore_graph( + workflow=workflow_snapshot( + make_workflow(workflow_id="workflow-1", tenant_id="tenant-1", app_id="app-1", graph="{}") + ), workflow_run_id=workflow_run_id, ) def test_validate_inputs_with_zero(): - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() var = VariableEntity( variable="test_var", @@ -106,7 +117,7 @@ def test_validate_inputs_with_zero(): def test_validate_input_with_none_for_required_variable(): - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() for var_type in VariableEntityType: var = VariableEntity( @@ -128,7 +139,7 @@ def test_validate_input_with_none_for_required_variable(): def test_validate_inputs_with_default_value(): """Test that default values are used when input is None for optional variables""" - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() # Test with string default value for TEXT_INPUT var_string = VariableEntity( @@ -343,7 +354,7 @@ def test_validate_inputs_with_default_value(): def test_validate_inputs_optional_file_with_empty_string(): """Test that optional FILE variable with empty string returns None""" - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() var_file = VariableEntity( variable="test_file", @@ -362,7 +373,7 @@ def test_validate_inputs_optional_file_with_empty_string(): def test_validate_inputs_optional_file_list_with_empty_list(): """Test that optional FILE_LIST variable with empty list returns empty list (not None)""" - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() var_file_list = VariableEntity( variable="test_file_list", @@ -383,7 +394,7 @@ def test_validate_inputs_optional_file_list_with_empty_list(): def test_validate_inputs_optional_file_list_with_empty_string(): """Test that optional FILE_LIST variable with empty string returns None""" - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() var_file_list = VariableEntity( variable="test_file_list", @@ -403,7 +414,7 @@ def test_validate_inputs_optional_file_list_with_empty_string(): def test_validate_inputs_required_file_with_empty_string_fails(): """Test that required FILE variable with empty string still fails validation""" - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() var_file = VariableEntity( variable="test_file", @@ -423,7 +434,7 @@ def test_validate_inputs_required_file_with_empty_string_fails(): def test_validate_inputs_optional_file_with_empty_string_ignores_default(): """Test that optional FILE variable with empty string returns None, not the default""" - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() var_file = VariableEntity( variable="test_file", @@ -442,16 +453,15 @@ def test_validate_inputs_optional_file_with_empty_string_ignores_default(): assert result is None -class TestBaseAppGeneratorExtras: +class TestAppInputAdapterExtras: 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( + managed_stream = consume_stream( response_stream(), worker_thread, ) @@ -465,7 +475,6 @@ def response_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 @@ -473,7 +482,7 @@ def response_stream(): yield {"event": "workflow_started"} yield {"event": "workflow_finished"} - managed_stream = base_app_generator._wrap_stream_with_worker_thread_join( + managed_stream = consume_stream( response_stream(), worker_thread, ) @@ -488,15 +497,15 @@ def test_join_worker_thread_warns_when_thread_remains_alive(self, caplog: pytest 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) + with caplog.at_level(logging.WARNING, logger="services.workflow.execution.generation_service"): + join_worker(worker_thread) worker_thread.join.assert_called_once_with(timeout=300) - assert "Possible app worker thread leak" in caplog.text + assert "did not stop" 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() + base_app_generator = AppInputAdapter() variables = [ VariableEntity( @@ -526,11 +535,11 @@ def test_prepare_user_inputs_converts_files_and_lists(self, monkeypatch: pytest. ] monkeypatch.setattr( - "core.app.apps.base_app_generator.file_factory.build_from_mapping", + "services.app.generation.input_adapter.file_factory.build_from_mapping", lambda mapping, tenant_id, config, strict_type_validation=False, access_controller=None: "file-object", ) monkeypatch.setattr( - "core.app.apps.base_app_generator.file_factory.build_from_mappings", + "services.app.generation.input_adapter.file_factory.build_from_mappings", lambda mappings, tenant_id, config, access_controller=None: ["file-1", "file-2"], ) @@ -551,7 +560,7 @@ def test_prepare_user_inputs_converts_files_and_lists(self, monkeypatch: pytest. assert prepared["json"] == {"key": "value"} def test_prepare_user_inputs_rejects_invalid_dict_inputs(self): - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() variables = [ VariableEntity( variable="text", @@ -569,7 +578,7 @@ def test_prepare_user_inputs_rejects_invalid_dict_inputs(self): ) def test_prepare_user_inputs_rejects_invalid_list_inputs(self): - base_app_generator = BaseAppGenerator() + base_app_generator = AppInputAdapter() variables = [ VariableEntity( variable="text", @@ -587,42 +596,42 @@ def test_prepare_user_inputs_rejects_invalid_list_inputs(self): ) def test_convert_to_event_stream(self): - base_app_generator = BaseAppGenerator() - - assert base_app_generator.convert_to_event_stream({"ok": True}) == {"ok": True} + assert convert_to_event_stream({"ok": True}) == {"ok": True} def _gen(): yield {"delta": "hi"} yield "ping" - converted = list(base_app_generator.convert_to_event_stream(_gen())) + converted = list(convert_to_event_stream(_gen())) assert converted[0].startswith("data: ") assert "\n\n" in converted[0] assert converted[1] == "event: ping\n\n" def test_get_draft_var_saver_factory_debugger(self): + from core.app.apps.draft_variable_saver import DraftVariableSaver, DraftVariableSaverFactory from models import Account - base_app_generator = BaseAppGenerator() + saver = create_autospec(DraftVariableSaver, instance=True, spec_set=True) + factory = create_autospec(DraftVariableSaverFactory, instance=True, spec_set=True) + factory.return_value = saver + provider = Mock(return_value=factory) + generator = AppInputAdapter(draft_variable_saver=provider) account = Account(name="Tester", email="tester@example.com") account.id = "account-id" - account.tenant_id = "tenant-id" - - factory = base_app_generator._get_draft_var_saver_factory( - InvokeFrom.DEBUGGER, - account, - tenant_id="tenant-id", + bound_factory = generator._get_draft_var_saver_factory(InvokeFrom.DEBUGGER, account, tenant_id="tenant-id") + assert ( + bound_factory( + app_id="app-id", node_id="node-id", node_type=BuiltinNodeTypes.START, node_execution_id="node-exec-id" + ) + is saver ) - saver = factory( - app_id="app-id", - node_id="node-id", - node_type=BuiltinNodeTypes.START, - node_execution_id="node-exec-id", + provider.assert_called_once_with("tenant-id", account) + assert factory.call_count == 1 + assert factory.call_args == call( + app_id="app-id", node_id="node-id", node_type=BuiltinNodeTypes.START, node_execution_id="node-exec-id" ) - assert saver is not None - DEFAULT_FILE_TENANT_ID = "tenant-1" DEFAULT_FILE_END_USER_ID = "end-user-1" @@ -674,7 +683,7 @@ def _file_input_variable(file: UploadFile, *, multiple: bool = False, required: def _prepare_file_inputs( variable: VariableEntity, inputs: dict[str, object], invoke_from: InvokeFrom = InvokeFrom.SERVICE_API ): - return BaseAppGenerator()._prepare_user_inputs( + return AppInputAdapter()._prepare_user_inputs( user_inputs=inputs, variables=[variable], tenant_id=DEFAULT_FILE_TENANT_ID, @@ -711,19 +720,6 @@ def test_omitted_file_uses_author_default( assert get_current_file_access_scope() is None -def test_app_input_adapter_uses_author_file_default(default_upload_file: UploadFile): - variable = _file_input_variable(default_upload_file) - with _file_access_scope(): - result = AppInputAdapter()._prepare_user_inputs( - user_inputs={}, - variables=[variable], - tenant_id=DEFAULT_FILE_TENANT_ID, - strict_type_validation=True, - )["file"] - assert isinstance(result, File) - assert resolve_file_record_id(result.reference) == default_upload_file.id - - @pytest.mark.parametrize("multiple", [False, True]) def test_webapp_can_submit_configured_default_explicitly(default_upload_file: UploadFile, multiple: bool): value = {"type": "image", "transfer_method": "local_file", "upload_file_id": default_upload_file.id} diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_runner.py b/api/tests/unit_tests/core/app/apps/test_base_app_runner.py index 563f8c188b62e4..0da00c9bda022b 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_runner.py @@ -2,12 +2,13 @@ import logging from dataclasses import dataclass +from decimal import Decimal from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest -from sqlalchemy import event, func, select -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy import func, select +from sqlalchemy.orm import Session from core.app.app_config.entities import ( AdvancedChatMessageEntity, @@ -15,7 +16,6 @@ AdvancedCompletionPromptTemplateEntity, PromptTemplateEntity, ) -from core.app.apps.base_app_runner import AppRunner from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import ( @@ -28,7 +28,6 @@ QueueAgentMessageEvent, QueueLLMChunkEvent, QueueMessageEndEvent, - QueueMessageFileEvent, ) from graphon.file import FileTransferMethod, FileType from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage @@ -40,9 +39,11 @@ ) from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelPropertyKey from graphon.model_runtime.errors.invoke import InvokeBadRequestError -from models.enums import CreatorUserRole -from models.model import App, AppMode, Message, MessageFile +from models.enums import ConversationFromSource, CreatorUserRole +from models.model import App, AppMode, MessageFile from models.tools import ToolFile +from services.app.generation.adapters.base_runner import AppRunner +from tests.helpers.unit.model_factories import make_app, make_conversation, make_message """Test multimodal image output handling in BaseAppRunner.""" @@ -50,6 +51,27 @@ class TestBaseAppRunnerMultimodal: """Test that BaseAppRunner correctly handles multimodal image content.""" + @pytest.fixture(autouse=True) + def message_owner(self, sqlite_session: Session, mock_tenant_id: str, mock_message_id: str) -> None: + sqlite_session.add_all( + [ + make_app(tenant_id=mock_tenant_id), + make_conversation(inputs={}, from_source=ConversationFromSource.API), + make_message( + message_id=mock_message_id, + inputs={}, + query="hi", + answer="", + message={}, + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + ), + ] + ) + sqlite_session.commit() + @pytest.fixture def mock_user_id(self): """Mock user ID.""" @@ -94,6 +116,8 @@ def test_handle_multimodal_image_content_with_url( mock_queue_manager, tool_file, sqlite_session: Session, + *, + app_records, ): """Test handling image from URL.""" # Arrange @@ -104,14 +128,13 @@ def test_handle_multimodal_image_content_with_url( mime_type="image/png", ) - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + with patch("services.app.generation.adapters.base_runner.ToolFileManager", autospec=True) as mock_mgr_class: # Setup mock tool file manager mock_mgr = MagicMock() mock_mgr.create_file_by_url.return_value = tool_file mock_mgr_class.return_value = mock_mgr - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, + message_file_id = AppRunner(records=app_records)._handle_multimodal_image_content( content=content, message_id=mock_message_id, user_id=mock_user_id, @@ -143,6 +166,8 @@ def test_handle_multimodal_image_content_with_base64( mock_queue_manager, tool_file, sqlite_session: Session, + *, + app_records, ): """Test handling image from base64 data.""" # Arrange @@ -158,14 +183,13 @@ def test_handle_multimodal_image_content_with_base64( mime_type="image/png", ) - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + with patch("services.app.generation.adapters.base_runner.ToolFileManager", autospec=True) as mock_mgr_class: # Setup mock tool file manager mock_mgr = MagicMock() mock_mgr.create_file_by_raw.return_value = tool_file mock_mgr_class.return_value = mock_mgr - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, + message_file_id = AppRunner(records=app_records)._handle_multimodal_image_content( content=content, message_id=mock_message_id, user_id=mock_user_id, @@ -193,6 +217,8 @@ def test_handle_multimodal_image_content_with_base64_data_uri( mock_queue_manager, tool_file, sqlite_session: Session, + *, + app_records, ): """Test handling image from base64 data with URI prefix.""" # Arrange @@ -206,14 +232,13 @@ def test_handle_multimodal_image_content_with_base64_data_uri( mime_type="image/png", ) - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + with patch("services.app.generation.adapters.base_runner.ToolFileManager", autospec=True) as mock_mgr_class: # Setup mock tool file manager mock_mgr = MagicMock() mock_mgr.create_file_by_raw.return_value = tool_file mock_mgr_class.return_value = mock_mgr - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, + message_file_id = AppRunner(records=app_records)._handle_multimodal_image_content( content=content, message_id=mock_message_id, user_id=mock_user_id, @@ -233,6 +258,8 @@ def test_handle_multimodal_image_content_without_url_or_base64( mock_message_id, mock_queue_manager, sqlite_session: Session, + *, + app_records, ): """Test handling image content without URL or base64 data.""" # Arrange @@ -243,9 +270,8 @@ def test_handle_multimodal_image_content_without_url_or_base64( mime_type="image/png", ) - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - result = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, + with patch("services.app.generation.adapters.base_runner.ToolFileManager", autospec=True) as mock_mgr_class: + result = AppRunner(records=app_records)._handle_multimodal_image_content( content=content, message_id=mock_message_id, user_id=mock_user_id, @@ -265,6 +291,8 @@ def test_handle_multimodal_image_content_with_error( mock_message_id, mock_queue_manager, sqlite_session: Session, + *, + app_records, ): """Test handling image content when an error occurs.""" # Arrange @@ -275,13 +303,12 @@ def test_handle_multimodal_image_content_with_error( mime_type="image/png", ) - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + with patch("services.app.generation.adapters.base_runner.ToolFileManager", autospec=True) as mock_mgr_class: mock_mgr = MagicMock() mock_mgr.create_file_by_url.side_effect = Exception("Network error") mock_mgr_class.return_value = mock_mgr - result = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, + result = AppRunner(records=app_records)._handle_multimodal_image_content( content=content, message_id=mock_message_id, user_id=mock_user_id, @@ -301,6 +328,8 @@ def test_handle_multimodal_image_content_debugger_mode( mock_queue_manager, tool_file, sqlite_session: Session, + *, + app_records, ): """Test that debugger mode sets correct created_by_role.""" # Arrange @@ -312,13 +341,12 @@ def test_handle_multimodal_image_content_debugger_mode( ) mock_queue_manager.invoke_from = InvokeFrom.DEBUGGER - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + with patch("services.app.generation.adapters.base_runner.ToolFileManager", autospec=True) as mock_mgr_class: mock_mgr = MagicMock() mock_mgr.create_file_by_url.return_value = tool_file mock_mgr_class.return_value = mock_mgr - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, + message_file_id = AppRunner(records=app_records)._handle_multimodal_image_content( content=content, message_id=mock_message_id, user_id=mock_user_id, @@ -338,6 +366,8 @@ def test_handle_multimodal_image_content_service_api_mode( mock_queue_manager, tool_file, sqlite_session: Session, + *, + app_records, ): """Test that service API mode sets correct created_by_role.""" # Arrange @@ -349,13 +379,12 @@ def test_handle_multimodal_image_content_service_api_mode( ) mock_queue_manager.invoke_from = InvokeFrom.SERVICE_API - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + with patch("services.app.generation.adapters.base_runner.ToolFileManager", autospec=True) as mock_mgr_class: mock_mgr = MagicMock() mock_mgr.create_file_by_url.return_value = tool_file mock_mgr_class.return_value = mock_mgr - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, + message_file_id = AppRunner(records=app_records)._handle_multimodal_image_content( content=content, message_id=mock_message_id, user_id=mock_user_id, @@ -417,8 +446,10 @@ def close(self) -> None: class TestAppRunner: - def test_recalc_llm_max_tokens_updates_parameters(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_recalc_llm_max_tokens_updates_parameters(self, monkeypatch: pytest.MonkeyPatch, *, app_records): + runner = AppRunner( + records=app_records, + ) model_schema = AIModelEntity.model_construct( model_properties={ModelPropertyKey.CONTEXT_SIZE: 100}, @@ -432,7 +463,7 @@ def test_recalc_llm_max_tokens_updates_parameters(self, monkeypatch: pytest.Monk ) monkeypatch.setattr( - "core.app.apps.base_app_runner.ModelInstance", + "services.app.generation.adapters.base_runner.ModelInstance", lambda provider_model_bundle, model: _TokenCountingModel(80), ) @@ -440,8 +471,12 @@ def test_recalc_llm_max_tokens_updates_parameters(self, monkeypatch: pytest.Monk assert model_config.parameters["max_tokens"] == 20 - def test_recalc_llm_max_tokens_returns_minus_one_when_no_context(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_recalc_llm_max_tokens_returns_minus_one_when_no_context( + self, monkeypatch: pytest.MonkeyPatch, *, app_records + ): + runner = AppRunner( + records=app_records, + ) model_schema = AIModelEntity.model_construct( model_properties={}, @@ -455,14 +490,16 @@ def test_recalc_llm_max_tokens_returns_minus_one_when_no_context(self, monkeypat ) monkeypatch.setattr( - "core.app.apps.base_app_runner.ModelInstance", + "services.app.generation.adapters.base_runner.ModelInstance", lambda provider_model_bundle, model: _TokenCountingModel(10), ) assert runner.recalc_llm_max_tokens(model_config, prompt_messages=[]) == -1 - def test_direct_output_streaming_publishes_chunks_and_end(self): - runner = AppRunner() + def test_direct_output_streaming_publishes_chunks_and_end(self, *, app_records): + runner = AppRunner( + records=app_records, + ) queue = _queue_manager() model_config = ModelConfigWithCredentialsEntity.model_construct(model="mock") app_generate_entity = EasyUIBasedAppGenerateEntity.model_construct(model_conf=model_config, stream=True) @@ -479,8 +516,10 @@ def test_direct_output_streaming_publishes_chunks_and_end(self): assert any(isinstance(event, QueueLLMChunkEvent) for event in events) assert isinstance(events[-1], QueueMessageEndEvent) - def test_handle_invoke_result_direct_publishes_end_event(self): - runner = AppRunner() + def test_handle_invoke_result_direct_publishes_end_event(self, *, app_records): + runner = AppRunner( + records=app_records, + ) queue = _queue_manager() llm_result = LLMResult( model="mock", @@ -497,8 +536,10 @@ def test_handle_invoke_result_direct_publishes_end_event(self): assert isinstance(_published_events(queue)[-1], QueueMessageEndEvent) - def test_handle_invoke_result_invalid_type_raises(self): - runner = AppRunner() + def test_handle_invoke_result_invalid_type_raises(self, *, app_records): + runner = AppRunner( + records=app_records, + ) queue = _queue_manager() with pytest.raises(NotImplementedError): @@ -508,8 +549,10 @@ def test_handle_invoke_result_invalid_type_raises(self): stream=True, ) - def test_organize_prompt_messages_simple_template(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_organize_prompt_messages_simple_template(self, monkeypatch: pytest.MonkeyPatch, *, app_records): + runner = AppRunner( + records=app_records, + ) model_config = ModelConfigWithCredentialsEntity.model_construct(mode="chat", stop=["STOP"]) prompt_template_entity = PromptTemplateEntity( prompt_type=PromptTemplateEntity.PromptType.SIMPLE, @@ -517,7 +560,7 @@ def test_organize_prompt_messages_simple_template(self, monkeypatch: pytest.Monk ) monkeypatch.setattr( - "core.app.apps.base_app_runner.SimplePromptTransform.get_prompt", + "services.app.generation.adapters.base_runner.SimplePromptTransform.get_prompt", lambda self, **kwargs: (["simple-message"], ["simple-stop"]), ) @@ -533,8 +576,12 @@ def test_organize_prompt_messages_simple_template(self, monkeypatch: pytest.Monk assert prompt_messages == ["simple-message"] assert stop == ["simple-stop"] - def test_organize_prompt_messages_advanced_completion_template(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_organize_prompt_messages_advanced_completion_template( + self, monkeypatch: pytest.MonkeyPatch, *, app_records + ): + runner = AppRunner( + records=app_records, + ) model_config = ModelConfigWithCredentialsEntity.model_construct(mode="completion", stop=[""]) captured: dict[str, object] = {} prompt_template_entity = PromptTemplateEntity( @@ -549,7 +596,9 @@ def _fake_advanced_prompt(self, **kwargs): captured.update(kwargs) return ["advanced-completion-message"] - monkeypatch.setattr("core.app.apps.base_app_runner.AdvancedPromptTransform.get_prompt", _fake_advanced_prompt) + monkeypatch.setattr( + "services.app.generation.adapters.base_runner.AdvancedPromptTransform.get_prompt", _fake_advanced_prompt + ) prompt_messages, stop = runner.organize_prompt_messages( app_record=App(mode=AppMode.CHAT.value), @@ -566,8 +615,10 @@ def _fake_advanced_prompt(self, **kwargs): assert memory_config.role_prefix.user == "U" assert memory_config.role_prefix.assistant == "A" - def test_organize_prompt_messages_advanced_chat_template(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_organize_prompt_messages_advanced_chat_template(self, monkeypatch: pytest.MonkeyPatch, *, app_records): + runner = AppRunner( + records=app_records, + ) model_config = ModelConfigWithCredentialsEntity.model_construct(mode="chat", stop=[""]) captured: dict[str, object] = {} prompt_template_entity = PromptTemplateEntity( @@ -584,7 +635,9 @@ def _fake_advanced_prompt(self, **kwargs): captured.update(kwargs) return ["advanced-chat-message"] - monkeypatch.setattr("core.app.apps.base_app_runner.AdvancedPromptTransform.get_prompt", _fake_advanced_prompt) + monkeypatch.setattr( + "services.app.generation.adapters.base_runner.AdvancedPromptTransform.get_prompt", _fake_advanced_prompt + ) prompt_messages, stop = runner.organize_prompt_messages( app_record=App(mode=AppMode.CHAT.value), @@ -599,8 +652,10 @@ def _fake_advanced_prompt(self, **kwargs): assert stop == [""] assert len(captured["prompt_template"]) == 2 - def test_organize_prompt_messages_advanced_missing_templates_raise(self): - runner = AppRunner() + def test_organize_prompt_messages_advanced_missing_templates_raise(self, *, app_records): + runner = AppRunner( + records=app_records, + ) with pytest.raises(InvokeBadRequestError, match="Advanced completion prompt template is required"): runner.organize_prompt_messages( @@ -620,8 +675,12 @@ def test_organize_prompt_messages_advanced_missing_templates_raise(self): files=[], ) - def test_handle_invoke_result_stream_routes_chunks_and_builds_message(self, caplog: pytest.LogCaptureFixture): - runner = AppRunner() + def test_handle_invoke_result_stream_routes_chunks_and_builds_message( + self, caplog: pytest.LogCaptureFixture, *, app_records + ): + runner = AppRunner( + records=app_records, + ) queue = _queue_manager() image_content = ImagePromptMessageContent( @@ -643,7 +702,7 @@ def _stream(): ), ) - with caplog.at_level(logging.WARNING, logger="core.app.apps.base_app_runner"): + with caplog.at_level(logging.WARNING, logger="services.app.generation.adapters.base_runner"): runner._handle_invoke_result( invoke_result=_stream(), queue_manager=queue, @@ -658,9 +717,11 @@ def _stream(): assert "Received multimodal output but missing required parameters" in caplog.messages def test_handle_invoke_result_stream_agent_mode_handles_multimodal_errors( - self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, *, app_records ): - runner = AppRunner() + runner = AppRunner( + records=app_records, + ) queue = _queue_manager() def raise_multimodal_error(**kwargs): @@ -693,7 +754,7 @@ def _stream(): ), ) - with caplog.at_level(logging.ERROR, logger="core.app.apps.base_app_runner"): + with caplog.at_level(logging.ERROR, logger="services.app.generation.adapters.base_runner"): runner._handle_invoke_result_stream( invoke_result=_stream(), queue_manager=queue, @@ -709,67 +770,12 @@ def _stream(): assert events[-1].llm_result.usage == usage assert "Failed to handle multimodal image output" in caplog.messages - def test_handle_invoke_result_stream_commits_message_file_before_publish( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session_factory: sessionmaker[Session], + def test_handle_invoke_result_stream_closes_generator_when_stopped( + self, monkeypatch: pytest.MonkeyPatch, *, app_records ): - runner = AppRunner() - monkeypatch.setattr( - runner, - "_handle_multimodal_image_content", - lambda **kwargs: "message-file-1", + runner = AppRunner( + records=app_records, ) - events: list[str] = [] - - def record_commit(session: Session) -> None: - events.append("commit") - - event.listen(sqlite_session_factory.class_, "after_commit", record_commit) - queue = _queue_manager() - original_publish = queue.publish - - def publish(event, pub_from): - if isinstance(event, QueueMessageFileEvent): - events.append("publish") - original_publish(event, pub_from) - - queue.publish = publish - - def stream(): - yield LLMResultChunk( - model="model", - prompt_messages=[AssistantPromptMessage(content="prompt")], - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage( - content=[ - ImagePromptMessageContent( - url="https://example.com/image.png", - format="png", - mime_type="image/png", - ) - ] - ), - ), - ) - - try: - runner._handle_invoke_result_stream( - invoke_result=stream(), - queue_manager=queue, - agent=False, - message_id="message-1", - user_id="user-1", - tenant_id="tenant-1", - ) - finally: - event.remove(sqlite_session_factory.class_, "after_commit", record_commit) - - assert events == ["commit", "publish"] - - def test_handle_invoke_result_stream_closes_generator_when_stopped(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() chunk = LLMResultChunk( model="stream-model", prompt_messages=[AssistantPromptMessage(content="prompt")], @@ -790,11 +796,10 @@ def test_handle_invoke_result_stream_closes_generator_when_stopped(self, monkeyp assert stream.closed is True @pytest.mark.parametrize("sqlite_session", [(MessageFile,)], indirect=True) - def test_handle_multimodal_image_content_fallback_return_branch( - self, - sqlite_session: Session, - ): - runner = AppRunner() + def test_handle_multimodal_image_content_fallback_return_branch(self, sqlite_session: Session, *, app_records): + runner = AppRunner( + records=app_records, + ) class _ToggleBool: def __init__(self, values: list[bool]): @@ -816,7 +821,6 @@ def __bool__(self): queue_manager = _queue_manager() runner._handle_multimodal_image_content( - session=sqlite_session, content=content, message_id="message-id", user_id="user-id", @@ -827,14 +831,16 @@ def __bool__(self): message_file_count = sqlite_session.scalar(select(func.count()).select_from(MessageFile)) assert message_file_count == 0 - def test_check_hosting_moderation_direct_output_called(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_check_hosting_moderation_direct_output_called(self, monkeypatch: pytest.MonkeyPatch, *, app_records): + runner = AppRunner( + records=app_records, + ) queue = _queue_manager() app_generate_entity = EasyUIBasedAppGenerateEntity.model_construct(stream=False) direct_output_calls: list[dict[str, object]] = [] monkeypatch.setattr( - "core.app.apps.base_app_runner.HostingModerationFeature.check", + "services.app.generation.adapters.base_runner.HostingModerationFeature.check", lambda self, application_generate_entity, prompt_messages: True, ) monkeypatch.setattr(runner, "direct_output", lambda **kwargs: direct_output_calls.append(kwargs)) @@ -848,10 +854,12 @@ def test_check_hosting_moderation_direct_output_called(self, monkeypatch: pytest assert result is True assert len(direct_output_calls) == 1 - def test_fill_in_inputs_from_external_data_tools(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_fill_in_inputs_from_external_data_tools(self, monkeypatch: pytest.MonkeyPatch, *, app_records): + runner = AppRunner( + records=app_records, + ) monkeypatch.setattr( - "core.app.apps.base_app_runner.ExternalDataFetch.fetch", + "services.app.generation.adapters.base_runner.ExternalDataFetch.fetch", lambda self, tenant_id, app_id, external_data_tools, inputs, query: {"foo": "bar"}, ) @@ -865,10 +873,12 @@ def test_fill_in_inputs_from_external_data_tools(self, monkeypatch: pytest.Monke assert result == {"foo": "bar"} - def test_moderation_for_inputs_returns_result(self, monkeypatch: pytest.MonkeyPatch): - runner = AppRunner() + def test_moderation_for_inputs_returns_result(self, monkeypatch: pytest.MonkeyPatch, *, app_records): + runner = AppRunner( + records=app_records, + ) monkeypatch.setattr( - "core.app.apps.base_app_runner.InputModeration.check", + "services.app.generation.adapters.base_runner.InputModeration.check", lambda self, app_id, tenant_id, app_config, inputs, query, message_id, trace_manager: (True, {}, ""), ) app_generate_entity = AppGenerateEntity.model_construct(app_config=None, trace_manager=None) @@ -883,26 +893,3 @@ def test_moderation_for_inputs_returns_result(self, monkeypatch: pytest.MonkeyPa ) assert result == (True, {}, "") - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_query_app_annotations_to_reply( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ): - runner = AppRunner() - monkeypatch.setattr( - "core.app.apps.base_app_runner.AnnotationReplyFeature.query", - lambda self, app_record, message, query, user_id, invoke_from, session: "reply", - ) - - response = runner.query_app_annotations_to_reply( - app_record=App(), - message=Message(), - query="hello", - user_id="user", - invoke_from=InvokeFrom.WEB_APP, - session=sqlite_session, - ) - - assert response == "reply" diff --git a/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py b/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py index 57bccb9696ba21..caf23705ba648f 100644 --- a/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py @@ -4,7 +4,6 @@ from types import SimpleNamespace import pytest -from sqlalchemy.orm import Session from core.app.app_config.entities import ( AppAdditionalFeatures, @@ -14,11 +13,11 @@ PromptTemplateEntity, ) from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.message_based_app_generator import MessageBasedAppGenerator from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom from models.model import App, AppMode, Conversation, Message +from services.app.generation.message_records import MessageBasedAppGenerator from services.errors.app_model_config import AppModelConfigBrokenError -from tests.helpers.unit.model_factories import make_app +from tests.helpers.unit.model_factories import make_app, make_end_user @dataclass(eq=False, repr=False) @@ -89,16 +88,20 @@ def _make_chat_generate_entity(app_config: EasyUIBasedAppConfig) -> ChatAppGener ) -def test_init_generate_records_skips_conversation_fields_for_non_conversation_entity(sqlite_session: Session): +def test_init_generate_records_skips_conversation_fields_for_non_conversation_entity( + app_records, *, annotation_replies +): app_config = _make_app_config(AppMode.COMPLETION) entity = DummyCompletionGenerateEntity(app_config=app_config) - generator = MessageBasedAppGenerator() + generator = MessageBasedAppGenerator( + annotations=annotation_replies, + records=app_records, + ) conversation, message = generator._init_generate_records( entity, conversation=None, - session=sqlite_session, ) assert conversation.id is not None @@ -107,16 +110,18 @@ def test_init_generate_records_skips_conversation_fields_for_non_conversation_en assert hasattr(entity, "is_new_conversation") is False -def test_init_generate_records_sets_conversation_fields_for_chat_entity(sqlite_session: Session): +def test_init_generate_records_sets_conversation_fields_for_chat_entity(app_records, *, annotation_replies): app_config = _make_app_config(AppMode.CHAT) entity = _make_chat_generate_entity(app_config) - generator = MessageBasedAppGenerator() + generator = MessageBasedAppGenerator( + annotations=annotation_replies, + records=app_records, + ) conversation, _ = generator._init_generate_records( entity, conversation=None, - session=sqlite_session, ) assert entity.conversation_id == conversation.id @@ -125,8 +130,13 @@ def test_init_generate_records_sets_conversation_fields_for_chat_entity(sqlite_s class TestMessageBasedAppGeneratorExtras: - def test_handle_response_closed_file_raises_stopped(self, monkeypatch: pytest.MonkeyPatch): - generator = MessageBasedAppGenerator() + def test_handle_response_closed_file_raises_stopped( + self, app_records, monkeypatch: pytest.MonkeyPatch, *, annotation_replies + ): + generator = MessageBasedAppGenerator( + annotations=annotation_replies, + records=app_records, + ) class _Pipeline: def __init__(self, **kwargs) -> None: @@ -136,7 +146,7 @@ def process(self): raise ValueError("I/O operation on closed file.") monkeypatch.setattr( - "core.app.apps.message_based_app_generator.EasyUIBasedGenerateTaskPipeline", + "services.app.generation.message_records.EasyUIBasedGenerateTaskPipeline", _Pipeline, ) @@ -146,31 +156,42 @@ def process(self): queue_manager=SimpleNamespace(), conversation=Conversation(id="conv", app_id="app"), message=Message(id="msg", app_id="app", conversation_id="conv"), + user=make_end_user(end_user_id="user-id"), stream=False, ) - def test_get_app_model_config_requires_valid_config(self, sqlite_session: Session): - generator = MessageBasedAppGenerator() + def test_get_app_model_config_requires_valid_config(self, app_records, *, annotation_replies): + generator = MessageBasedAppGenerator( + annotations=annotation_replies, + records=app_records, + ) app_model = _app() - session = sqlite_session with pytest.raises(AppModelConfigBrokenError): - generator._get_app_model_config(app_model, conversation=None, session=session) + generator._get_app_model_config(app_model, conversation=None) conversation = Conversation(id="conversation-id", app_id="app", app_model_config_id="missing-id") with pytest.raises(AppModelConfigBrokenError): generator._get_app_model_config( app_model=_app(), conversation=conversation, - session=session, ) - def test_get_conversation_introduction_handles_missing_inputs(self): + def test_get_conversation_introduction_handles_missing_inputs(self, app_records, *, annotation_replies): app_config = _make_app_config(AppMode.CHAT) app_config.additional_features.opening_statement = "Hello {{name}}" entity = _make_chat_generate_entity(app_config) entity.inputs = {} - generator = MessageBasedAppGenerator() + generator = MessageBasedAppGenerator( + annotations=annotation_replies, + records=app_records, + ) assert generator._get_conversation_introduction(entity) == "Hello {name}" + + +@pytest.fixture(autouse=True) +def message_app(sqlite_session): + sqlite_session.add(make_app(app_id="app-id", tenant_id="tenant-id")) + sqlite_session.commit() diff --git a/api/tests/unit_tests/core/app/apps/test_message_generator.py b/api/tests/unit_tests/core/app/apps/test_message_generator.py index a7ecf79b1e0c6c..6bdeb06a54a07b 100644 --- a/api/tests/unit_tests/core/app/apps/test_message_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_message_generator.py @@ -1,9 +1,9 @@ from unittest.mock import Mock, patch -from core.app.apps.message_generator import MessageGenerator from core.app.entities.task_entities import StreamEvent from libs.broadcast_channel.channel import SupportsPreparedSubscription from models.model import AppMode +from services.workflow.execution.adapters.response_stream import WorkflowEventStream class _PreparedSubscriber(SupportsPreparedSubscription): @@ -29,16 +29,18 @@ def subscribe(self): raise AssertionError("event retrieval must use the topic's subscriber view") -class TestMessageGenerator: +class TestWorkflowEventStream: def test_get_response_topic(self): channel = Mock() channel.topic.return_value = "topic" - with patch("core.app.apps.message_generator.get_pubsub_broadcast_channel", return_value=channel): - topic = MessageGenerator.get_response_topic(AppMode.WORKFLOW, "run-1") + with patch( + "services.workflow.execution.adapters.response_stream.get_pubsub_broadcast_channel", return_value=channel + ): + topic = WorkflowEventStream.get_response_topic(AppMode.WORKFLOW, "run-1") assert topic == "topic" - expected_key = MessageGenerator._make_channel_key(AppMode.WORKFLOW, "run-1") + expected_key = WorkflowEventStream._make_channel_key(AppMode.WORKFLOW, "run-1") channel.topic.assert_called_once_with(expected_key) def test_retrieve_events_passes_arguments(self): @@ -46,13 +48,17 @@ def test_retrieve_events_passes_arguments(self): topic.as_subscriber.return_value = topic topic.subscribe.return_value = "subscription" with ( - patch("core.app.apps.message_generator.MessageGenerator.get_response_topic", return_value=topic), patch( - "core.app.apps.message_generator.stream_topic_events", return_value=iter([{"event": "ping"}]) + "services.workflow.execution.adapters.response_stream.WorkflowEventStream.get_response_topic", + return_value=topic, + ), + patch( + "services.workflow.execution.adapters.response_stream.stream_topic_events", + return_value=iter([{"event": "ping"}]), ) as mock_stream, ): events = list( - MessageGenerator.retrieve_events( + WorkflowEventStream.retrieve_events( AppMode.WORKFLOW, "run-1", idle_timeout=1, @@ -75,10 +81,15 @@ def test_retrieve_events_passes_arguments(self): def test_retrieve_events_uses_prepared_subscription_capability(self): topic = _TopicWithSeparateSubscriberView() with ( - patch("core.app.apps.message_generator.MessageGenerator.get_response_topic", return_value=topic), - patch("core.app.apps.message_generator.stream_topic_events", return_value=iter([])) as mock_stream, + patch( + "services.workflow.execution.adapters.response_stream.WorkflowEventStream.get_response_topic", + return_value=topic, + ), + patch( + "services.workflow.execution.adapters.response_stream.stream_topic_events", return_value=iter([]) + ) as mock_stream, ): - events = MessageGenerator.retrieve_events(AppMode.WORKFLOW, "run-1") + events = WorkflowEventStream.retrieve_events(AppMode.WORKFLOW, "run-1") assert topic.subscriber.prepare_calls == 1 mock_stream.assert_called_once_with( diff --git a/api/tests/unit_tests/core/app/apps/test_streaming_utils.py b/api/tests/unit_tests/core/app/apps/test_streaming_utils.py index 66afebb3e60f4a..3fb75e6e5b1378 100644 --- a/api/tests/unit_tests/core/app/apps/test_streaming_utils.py +++ b/api/tests/unit_tests/core/app/apps/test_streaming_utils.py @@ -5,11 +5,11 @@ import pytest -from core.app.apps.message_based_app_generator import MessageBasedAppGenerator from core.app.apps.streaming_utils import _normalize_terminal_events, stream_topic_events from core.app.entities.task_entities import StreamEvent from libs.broadcast_channel.channel import SupportsPreparedSubscription from models.model import AppMode +from services.workflow.execution.adapters.response_stream import WorkflowEventStream class FakeSubscription: @@ -107,14 +107,14 @@ def test_retrieve_events_falls_back_to_subscribe_and_invokes_hook_after_entry(mo def fake_get_response_topic(cls, app_mode, workflow_run_id): return topic - monkeypatch.setattr(MessageBasedAppGenerator, "get_response_topic", classmethod(fake_get_response_topic)) + monkeypatch.setattr(WorkflowEventStream, "get_response_topic", classmethod(fake_get_response_topic)) def on_subscribe() -> None: assert topic.subscribed is True event = {"event": StreamEvent.WORKFLOW_FINISHED.value} topic.publish(json.dumps(event).encode()) - generator = MessageBasedAppGenerator.retrieve_events( + generator = WorkflowEventStream.retrieve_events( AppMode.WORKFLOW, "workflow-run-id", idle_timeout=0.5, @@ -137,9 +137,9 @@ def test_retrieve_events_prepares_capable_topic_before_generator_iteration(monke def fake_get_response_topic(cls, app_mode, workflow_run_id): return topic - monkeypatch.setattr(MessageBasedAppGenerator, "get_response_topic", classmethod(fake_get_response_topic)) + monkeypatch.setattr(WorkflowEventStream, "get_response_topic", classmethod(fake_get_response_topic)) - generator = MessageBasedAppGenerator.retrieve_events( + generator = WorkflowEventStream.retrieve_events( AppMode.WORKFLOW, "workflow-run-id", idle_timeout=0.5, @@ -161,10 +161,10 @@ def test_retrieve_events_propagates_preparation_error_before_generator_iteration def fake_get_response_topic(cls, app_mode, workflow_run_id): return topic - monkeypatch.setattr(MessageBasedAppGenerator, "get_response_topic", classmethod(fake_get_response_topic)) + monkeypatch.setattr(WorkflowEventStream, "get_response_topic", classmethod(fake_get_response_topic)) with pytest.raises(RuntimeError, match="prepare failed"): - MessageBasedAppGenerator.retrieve_events( + WorkflowEventStream.retrieve_events( AppMode.WORKFLOW, "workflow-run-id", idle_timeout=0.5, diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py index b65aa1964249a2..a4d9f2f81b257f 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py @@ -4,27 +4,30 @@ from collections.abc import Generator, Mapping from copy import deepcopy from datetime import UTC, datetime, timedelta +from decimal import Decimal from time import perf_counter from types import SimpleNamespace from typing import override -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import MagicMock, patch from urllib.parse import parse_qs, urlparse import pytest from pydantic import ValidationError from pytest_mock import MockerFixture from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from configs import dify_config from core.app.app_config.entities import WorkflowUIBasedAppConfig from core.app.apps.base_app_queue_manager import PublishFrom -from core.app.apps.common import workflow_response_converter from core.app.apps.execution_coordinator import AppExecutionState -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow.app_runner import WorkflowAppRunner -from core.app.apps.workflow_app_runner import PreparedWorkflowRun, WorkflowBasedAppRunner -from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom, WorkflowAppGenerateEntity +from core.app.entities.app_invoke_entities import ( + DIFY_RUN_CONTEXT_KEY, + AdvancedChatAppGenerateEntity, + InvokeFrom, + UserFrom, + WorkflowAppGenerateEntity, +) from core.app.entities.queue_entities import ( AppQueueEvent, QueueAgentLogEvent, @@ -47,18 +50,11 @@ QueueWorkflowSucceededEvent, ) from core.credit_usage import CreditUsageAppType -from core.db.session_factory import session_factory -from core.repositories.human_input_repository import HumanInputFormRepositoryImpl, HumanInputFormSubmissionRepository -from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.node_runtime import DifyFileReferenceFactory from core.workflow.nodes.agent.events import NodeRunAgentLogEvent -from core.workflow.nodes.human_input.boundary import human_input_container_selector -from core.workflow.nodes.human_input.callback import DifyHITLCallback from core.workflow.nodes.human_input.pause_reason import HumanInputRequired from core.workflow.system_variables import default_system_variables -from core.workflow.workflow_entry import WorkflowEntry, iter_dify_graph_engine_events from enums.human_input import HumanInputFormKind, HumanInputFormStatus +from extensions.application_services.workflow import WorkflowExecutionDependencies from extensions.storage.storage_type import StorageType from graphon.engine.layer import Layer from graphon.engine.ready_queue import StartTask @@ -71,6 +67,7 @@ GraphRunSucceededEvent, NodeRunExceptionEvent, NodeRunFailedEvent, + NodeRunHumanInputFormFilledEvent, NodeRunIterationStartedEvent, NodeRunIterationSucceededEvent, NodeRunLoopFailedEvent, @@ -92,6 +89,8 @@ from graphon.nodes.human_input.human_input_node import HumanInputNode from graphon.runtime import RuntimeState, VariablePool from graphon.runtime.execution import ROOT_FRAME_ID +from graphon.variable_loader import DUMMY_VARIABLE_LOADER +from graphon.variables.factory import build_segment from graphon.variables.variables import StringVariable from models import Account, UploadFile from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom @@ -105,16 +104,26 @@ ParagraphInputConfig, UserActionConfig, ) -from models.model import AppMode +from models.model import AppMode, ConversationFromSource from models.workflow import Workflow, WorkflowKind, WorkflowNodeExecutionTriggeredFrom from models.workflow import WorkflowType as ModelWorkflowType +from repositories.human_input.form_repository import HumanInputFormSubmissionRepository +from repositories.workflow.definition_repository import workflow_snapshot +from services.app.generation.ports import MessageIdentity +from services.tools.provider_queries import ToolProviders +from services.workflow.execution.adapters.events import WorkflowEventPublisher, enqueue_human_input_notifications +from services.workflow.execution.adapters.graph import WorkflowGraphBuilder +from services.workflow.execution.adapters.human_input import DifyHITLCallback +from services.workflow.execution.adapters.node_runtime import DifyFileReferenceFactory +from services.workflow.execution.adapters.workflow.app_queue_manager import WorkflowAppQueueManager +from services.workflow.execution.adapters.workflow.app_runner import WorkflowAppRunner +from services.workflow.execution.adapters.workflow_entry import WorkflowEntry, iter_dify_graph_engine_events +from services.workflow.execution.contracts import PreparedWorkflowRun +from services.workflow.execution.ports import ExecutionContexts, WorkflowRuntime from services.workflow_run_agg import WorkflowRunAgg from services.workflow_run_index import WorkflowRunIndex -from tests.helpers.unit.model_factories import make_workflow -from tests.unit_tests.core.app.apps.advanced_chat.test_generate_task_pipeline import ( - _build_generate_entity, - _build_pipeline, -) +from tests.helpers.unit.model_factories import make_app, make_conversation, make_message, make_workflow +from tests.unit_tests.core.app.apps.advanced_chat.test_generate_task_pipeline import _build_pipeline from tests.unit_tests.core.app.apps.common.test_workflow_response_converter import ( _workflow_response_converter_human_input_build_converter as _build_converter, ) @@ -130,9 +139,9 @@ def _workflow(graph: dict[str, object] | None = None) -> Workflow: ) -def _workflow_entry(runner: WorkflowBasedAppRunner, state: RuntimeState | None = None) -> WorkflowEntry: +def _workflow_entry(_runner: object | None = None, state: RuntimeState | None = None) -> WorkflowEntry: state = state or RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0.0) - graph = runner._init_graph( + graph = WorkflowGraphBuilder(app_id="app").build( graph_config={ "nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], "edges": [], @@ -142,7 +151,13 @@ def _workflow_entry(runner: WorkflowBasedAppRunner, state: RuntimeState | None = invoke_from=InvokeFrom.DEBUGGER, call_depth=0, ) - return WorkflowEntry(graph_config=graph.graph_config, graph=graph, call_depth=0, graph_runtime_state=state) + return WorkflowEntry( + graph_config=graph.graph_config, + graph=graph, + call_depth=0, + graph_runtime_state=state, + human_form_reader=MagicMock(), + ) def _nested_container_graph(ownership: str) -> dict: @@ -211,820 +226,846 @@ def _nested_container_graph(ownership: str) -> dict: } -class TestWorkflowAppRunnerCoreCases: - class TestWorkflowBasedAppRunner: - @pytest.mark.parametrize( - ("event_class", "node_type"), - [ - (NodeRunIterationStartedEvent, BuiltinNodeTypes.ITERATION), - (NodeRunIterationSucceededEvent, BuiltinNodeTypes.ITERATION), - (NodeRunLoopStartedEvent, BuiltinNodeTypes.LOOP), - (NodeRunLoopSucceededEvent, BuiltinNodeTypes.LOOP), - ], +class TestWorkflowExecutionAdapters: + @pytest.mark.parametrize( + ("event_class", "node_type"), + [ + (NodeRunIterationStartedEvent, BuiltinNodeTypes.ITERATION), + (NodeRunIterationSucceededEvent, BuiltinNodeTypes.ITERATION), + (NodeRunLoopStartedEvent, BuiltinNodeTypes.LOOP), + (NodeRunLoopSucceededEvent, BuiltinNodeTypes.LOOP), + ], + ) + def test_container_queue_events_preserve_lifecycle_and_ancestor_metadata( + self, + event_class: type[ + NodeRunIterationStartedEvent + | NodeRunIterationSucceededEvent + | NodeRunLoopStartedEvent + | NodeRunLoopSucceededEvent + ], + node_type: str, + ) -> None: + queue = WorkflowAppQueueManager( + task_id="metadata-contract", user_id="account", invoke_from=InvokeFrom.DEBUGGER, app_mode="workflow" ) - def test_container_queue_events_preserve_lifecycle_and_ancestor_metadata( - self, - event_class: type[ - NodeRunIterationStartedEvent - | NodeRunIterationSucceededEvent - | NodeRunLoopStartedEvent - | NodeRunLoopSucceededEvent - ], - node_type: str, - ) -> None: - queue = WorkflowAppQueueManager( - task_id="metadata-contract", user_id="account", invoke_from=InvokeFrom.DEBUGGER, app_mode="workflow" - ) - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _workflow_entry(runner) - event = event_class( - id="container-execution", - node_id="container", - node_type=node_type, - node_title="Container", - start_at=datetime.now(UTC), - metadata={"iteration_length": 2, "total_tokens": 1}, - node_run_result=NodeRunResult( - metadata={ - WorkflowNodeExecutionMetadataKey.LOOP_ID: "outer", - WorkflowNodeExecutionMetadataKey.LOOP_INDEX: 3, - WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 2, - } - ), - ) - original_event = event.model_dump() + publisher = WorkflowEventPublisher(queue_manager=queue, notify_pause=enqueue_human_input_notifications) + entry = _workflow_entry() + event = event_class( + id="container-execution", + node_id="container", + node_type=node_type, + node_title="Container", + start_at=datetime.now(UTC), + metadata={"iteration_length": 2, "total_tokens": 1}, + node_run_result=NodeRunResult( + metadata={ + WorkflowNodeExecutionMetadataKey.LOOP_ID: "outer", + WorkflowNodeExecutionMetadataKey.LOOP_INDEX: 3, + WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 2, + } + ), + ) + original_event = event.model_dump() - runner.handle_event(entry, event) - queue.stop_listen(execution_state=AppExecutionState.TERMINAL) - messages = list(queue.listen()) + publisher.publish(entry, event) + queue.stop_listen(execution_state=AppExecutionState.TERMINAL) + messages = list(queue.listen()) - assert len(messages) == 1 - queued_event = messages[0].event - assert isinstance( - queued_event, - QueueIterationStartEvent | QueueIterationCompletedEvent | QueueLoopStartEvent | QueueLoopCompletedEvent, - ) - assert queued_event.metadata == { - "iteration_length": 2, - "total_tokens": 2, - "loop_id": "outer", - "loop_index": 3, - } - assert event.model_dump() == original_event + assert len(messages) == 1 + queued_event = messages[0].event + assert isinstance( + queued_event, + QueueIterationStartEvent | QueueIterationCompletedEvent | QueueLoopStartEvent | QueueLoopCompletedEvent, + ) + assert queued_event.metadata == { + "iteration_length": 2, + "total_tokens": 2, + "loop_id": "outer", + "loop_index": 3, + } + assert event.model_dump() == original_event - def test_resolve_user_from(self): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + def test_resolve_user_from(self): + runner = WorkflowGraphBuilder(app_id="app") - assert runner._resolve_user_from(InvokeFrom.EXPLORE) == UserFrom.ACCOUNT - assert runner._resolve_user_from(InvokeFrom.DEBUGGER) == UserFrom.ACCOUNT - assert runner._resolve_user_from(InvokeFrom.WEB_APP) == UserFrom.END_USER + assert runner.resolve_user_from(InvokeFrom.EXPLORE) == UserFrom.ACCOUNT + assert runner.resolve_user_from(InvokeFrom.DEBUGGER) == UserFrom.ACCOUNT + assert runner.resolve_user_from(InvokeFrom.WEB_APP) == UserFrom.END_USER - def test_init_graph_validates_graph_structure(self): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + def test_init_graph_validates_graph_structure(self): + runner = WorkflowGraphBuilder(app_id="app") - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) - with pytest.raises(ValidationError, match="nodes"): - runner._init_graph( - graph_config={}, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - with pytest.raises(ValidationError, match="nodes"): - runner._init_graph( - graph_config={"nodes": {}, "edges": []}, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - with pytest.raises(ValidationError, match="edges"): - runner._init_graph( - graph_config={"nodes": [], "edges": {}}, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - def test_init_graph_includes_trace_session_id_in_run_context(self, monkeypatch: pytest.MonkeyPatch): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - captured = {} - - def fake_from_graph_init_context(**kwargs): - graph_init_context = kwargs["graph_init_context"] - captured["run_context"] = graph_init_context.run_context - captured["call_depth"] = graph_init_context.call_depth - captured["runtime_state"] = kwargs["runtime_state"] - return SimpleNamespace() - - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.DifyNodeFactory.from_graph_init_context", - fake_from_graph_init_context, + with pytest.raises(ValidationError, match="nodes"): + runner.build( + graph_config={}, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", lambda **_kwargs: SimpleNamespace()) - runner._init_graph( - graph_config={"nodes": [], "edges": []}, + with pytest.raises(ValidationError, match="nodes"): + runner.build( + graph_config={"nodes": {}, "edges": []}, graph_runtime_state=runtime_state, user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - root_node_id="root", - trace_session_id="session-1", - call_depth=3, + call_depth=0, ) - assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" - assert captured["call_depth"] == 3 - assert captured["runtime_state"] is runtime_state - - @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) - def test_init_graph_preserves_container_execution_scopes(self, ownership: str) -> None: - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - graph_config = _nested_container_graph(ownership) - original_config = deepcopy(graph_config) - graph = runner._init_graph( - graph_config=graph_config, + with pytest.raises(ValidationError, match="edges"): + runner.build( + graph_config={"nodes": [], "edges": {}}, graph_runtime_state=runtime_state, user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - root_node_id="start", call_depth=0, ) - assert set(graph.nodes) == {"start", "outer", "end"} - outer_graph = Graph.init( - graph_config=graph.graph_config, - node_factory=graph.node_factory, - root_node_id="outer-start", - container_id="outer", - ) - assert set(outer_graph.nodes) == {"outer-start", "inner"} - inner_graph = Graph.init( - graph_config=outer_graph.graph_config, - node_factory=outer_graph.node_factory, - root_node_id="inner-start", - container_id="inner", - ) - assert set(inner_graph.nodes) == {"inner-start", "deep"} - deep_graph = Graph.init( - graph_config=inner_graph.graph_config, - node_factory=inner_graph.node_factory, - root_node_id="deep-start", - container_id="deep", - ) - assert set(deep_graph.nodes) == {"deep-start"} - assert graph_config == original_config - - def test_prepare_single_node_execution_requires_run(self): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - - workflow = _workflow() - - with pytest.raises(ValueError, match="Neither single_iteration_run nor single_loop_run"): - runner._prepare_single_node_execution( - workflow, None, None, user_id="00000000-0000-0000-0000-000000000001" - ) - - @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) - @pytest.mark.parametrize("node_id", ["outer", "inner", "deep"]) - def test_single_container_debugging_preserves_descendants(self, ownership: str, node_id: str): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) + def test_init_graph_includes_trace_session_id_in_run_context(self, monkeypatch: pytest.MonkeyPatch): + runner = WorkflowGraphBuilder(app_id="app") + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + captured = {} - graph_config = _nested_container_graph(ownership) - workflow = _workflow(graph_config) - workflow.id = "workflow" - graph_runtime_state.variable_pool.add(["start", "items"], ["first", "second"]) - - graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id=node_id, - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="iteration", - user_id="00000000-0000-0000-0000-000000000001", - ) + def fake_from_graph_init_context(**kwargs): + graph_init_context = kwargs["graph_init_context"] + captured["run_context"] = graph_init_context.run_context + captured["call_depth"] = graph_init_context.call_depth + captured["runtime_state"] = kwargs["runtime_state"] + return SimpleNamespace() - assert set(graph.nodes) == {node_id} - assert variable_pool is graph_runtime_state.variable_pool - deep_graph = Graph.init( - graph_config=graph.graph_config, - node_factory=graph.node_factory, - root_node_id="deep-start", - container_id="deep", - ) - assert set(deep_graph.nodes) == {"deep-start"} - - def test_get_graph_and_variable_pool_for_single_node_run_does_not_mutate_inputs( - self, monkeypatch: pytest.MonkeyPatch - ): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.DifyNodeFactory.from_graph_init_context", + fake_from_graph_init_context, + ) + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.Graph.init", lambda **_kwargs: SimpleNamespace() + ) + + runner.build( + graph_config={"nodes": [], "edges": []}, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + root_node_id="root", + trace_session_id="session-1", + call_depth=3, + ) + + assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" + assert captured["call_depth"] == 3 + assert captured["runtime_state"] is runtime_state + + @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) + def test_init_graph_preserves_container_execution_scopes(self, ownership: str) -> None: + runner = WorkflowGraphBuilder(app_id="app") + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + graph_config = _nested_container_graph(ownership) + original_config = deepcopy(graph_config) + graph = runner.build( + graph_config=graph_config, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + root_node_id="start", + call_depth=0, + ) + + assert set(graph.nodes) == {"start", "outer", "end"} + outer_graph = Graph.init( + graph_config=graph.graph_config, + node_factory=graph.node_factory, + root_node_id="outer-start", + container_id="outer", + ) + assert set(outer_graph.nodes) == {"outer-start", "inner"} + inner_graph = Graph.init( + graph_config=outer_graph.graph_config, + node_factory=outer_graph.node_factory, + root_node_id="inner-start", + container_id="inner", + ) + assert set(inner_graph.nodes) == {"inner-start", "deep"} + deep_graph = Graph.init( + graph_config=inner_graph.graph_config, + node_factory=inner_graph.node_factory, + root_node_id="deep-start", + container_id="deep", + ) + assert set(deep_graph.nodes) == {"deep-start"} + assert graph_config == original_config + + def test_prepare_single_node_execution_requires_run(self): + runner = WorkflowGraphBuilder(app_id="app") + + workflow = _workflow() + + with pytest.raises(ValueError, match="Neither single_iteration_run nor single_loop_run"): + runner.build_single_node(workflow, None, None, user_id="00000000-0000-0000-0000-000000000001") + + @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) + @pytest.mark.parametrize("node_id", ["outer", "inner", "deep"]) + def test_single_container_debugging_preserves_descendants(self, ownership: str, node_id: str): + runner = WorkflowGraphBuilder(app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + + graph_config = _nested_container_graph(ownership) + workflow = _workflow(graph_config) + workflow.id = "workflow" + graph_runtime_state.variable_pool.add(["start", "items"], ["first", "second"]) + + graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id=node_id, + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="iteration", + user_id="00000000-0000-0000-0000-000000000001", + ) + + assert set(graph.nodes) == {node_id} + assert variable_pool is graph_runtime_state.variable_pool + deep_graph = Graph.init( + graph_config=graph.graph_config, + node_factory=graph.node_factory, + root_node_id="deep-start", + container_id="deep", + ) + assert set(deep_graph.nodes) == {"deep-start"} + + def test_get_graph_and_variable_pool_for_single_node_run_does_not_mutate_inputs( + self, monkeypatch: pytest.MonkeyPatch + ): + runner = WorkflowGraphBuilder(app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + graph_config = { + "nodes": [ + {"id": "node-1", "data": {"type": "start", "version": "1"}}, + {"id": "node-child", "data": {"type": "start", "version": "1", "iteration_id": "node-1"}}, + {"id": "outside", "data": {"type": "start", "version": "1"}}, + ], + "edges": [ + {"id": "edge-1", "source": "node-1", "target": "node-child"}, + {"id": "edge-2", "source": "outside", "target": "node-1"}, + ], + } + expected_graph_config = deepcopy(graph_config) + workflow = SimpleNamespace(tenant_id="tenant", id="workflow", graph_dict=graph_config) + user_inputs = {"question": "hello"} + expected_user_inputs = dict(user_inputs) + captured_graph_config = {} + + def fake_graph_init(**kwargs): + captured_graph_config["value"] = kwargs["graph_config"] + return SimpleNamespace() + + class _NodeCls: + @staticmethod + def extract_variable_selector_to_variable_mapping(graph_config, config): + return {} + + from services.workflow.execution.adapters import graph as workflow_app_runner + + monkeypatch.setattr("services.workflow.execution.adapters.graph.Graph.init", fake_graph_init) + monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.load_into_variable_pool", + lambda **kwargs: kwargs["user_inputs"].update({"from_loader": "local-only"}), + ) + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.WorkflowEntry.mapping_user_inputs_to_variable_pool", + lambda **kwargs: kwargs["user_inputs"].update({"from_mapping": "local-only"}), + ) + + graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="node-1", + user_inputs=user_inputs, + graph_runtime_state=graph_runtime_state, + node_type_label="iteration", + user_id="00000000-0000-0000-0000-000000000001", + ) + + assert graph is not None + assert variable_pool is graph_runtime_state.variable_pool + assert workflow.graph_dict == expected_graph_config + assert user_inputs == expected_user_inputs + assert captured_graph_config["value"] is not workflow.graph_dict + assert [node["id"] for node in captured_graph_config["value"]["nodes"]] == ["node-1", "node-child"] + assert captured_graph_config["value"]["edges"] == [{"id": "edge-1", "source": "node-1", "target": "node-child"}] + + def test_get_graph_and_variable_pool_for_single_node_run_includes_trace_session_id( + self, monkeypatch: pytest.MonkeyPatch + ): + runner = WorkflowGraphBuilder(app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + graph_config = { + "nodes": [{"id": "node-1", "data": {"type": "start", "version": "1"}}], + "edges": [], + } + workflow = _workflow(graph_config) + workflow.id = "workflow" + captured = {} + + def fake_from_graph_init_context(**kwargs): + captured["run_context"] = kwargs["graph_init_context"].run_context + captured["runtime_state"] = kwargs["runtime_state"] + return SimpleNamespace() + + class _NodeCls: + @staticmethod + def extract_variable_selector_to_variable_mapping(graph_config, config): + return {} + + from services.workflow.execution.adapters import graph as workflow_app_runner + + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.DifyNodeFactory.from_graph_init_context", + fake_from_graph_init_context, + ) + monkeypatch.setattr("services.workflow.execution.adapters.graph.Graph.init", lambda **kwargs: SimpleNamespace()) + monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) + monkeypatch.setattr("services.workflow.execution.adapters.graph.load_into_variable_pool", lambda **kwargs: None) + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.WorkflowEntry.mapping_user_inputs_to_variable_pool", + lambda **kwargs: None, + ) + + runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="node-1", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="iteration", + user_id="00000000-0000-0000-0000-000000000001", + trace_session_id="session-1", + ) + + assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" + assert captured["runtime_state"] is graph_runtime_state + + def test_get_graph_and_variable_pool_preloads_constructor_variables_before_graph_init( + self, monkeypatch: pytest.MonkeyPatch + ): + variable_loader = SimpleNamespace( + load_variables=lambda selectors: ( + [ + StringVariable( + name="conversation_id", + value="conv-1", + selector=["sys", "conversation_id"], + ) + ] + if selectors + else [] ) - graph_config = { + ) + runner = WorkflowGraphBuilder(variable_loader=variable_loader, app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + + workflow = _workflow( + { "nodes": [ - {"id": "node-1", "data": {"type": "start", "version": "1"}}, - {"id": "node-child", "data": {"type": "start", "version": "1", "iteration_id": "node-1"}}, - {"id": "outside", "data": {"type": "start", "version": "1"}}, - ], - "edges": [ - {"id": "edge-1", "source": "node-1", "target": "node-child"}, - {"id": "edge-2", "source": "outside", "target": "node-1"}, + {"id": "loop-node", "data": {"type": "loop", "version": "1", "title": "Loop"}}, + { + "id": "llm-child", + "data": { + "type": "llm", + "version": "1", + "loop_id": "loop-node", + "memory": object(), + }, + }, ], + "edges": [], } - expected_graph_config = deepcopy(graph_config) - workflow = SimpleNamespace(tenant_id="tenant", id="workflow", graph_dict=graph_config) - user_inputs = {"question": "hello"} - expected_user_inputs = dict(user_inputs) - captured_graph_config = {} - - def fake_graph_init(**kwargs): - captured_graph_config["value"] = kwargs["graph_config"] - return SimpleNamespace() - - class _NodeCls: - @staticmethod - def extract_variable_selector_to_variable_mapping(graph_config, config): - return {} - - from core.app.apps import workflow_app_runner - - monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", fake_graph_init) - monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.load_into_variable_pool", - lambda **kwargs: kwargs["user_inputs"].update({"from_loader": "local-only"}), - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", - lambda **kwargs: kwargs["user_inputs"].update({"from_mapping": "local-only"}), - ) + ) + workflow.id = "workflow" - graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="node-1", - user_inputs=user_inputs, - graph_runtime_state=graph_runtime_state, - node_type_label="iteration", - user_id="00000000-0000-0000-0000-000000000001", - ) + class _LoopNodeCls: + @staticmethod + def extract_variable_selector_to_variable_mapping(graph_config, config): + return {} - assert graph is not None - assert variable_pool is graph_runtime_state.variable_pool - assert workflow.graph_dict == expected_graph_config - assert user_inputs == expected_user_inputs - assert captured_graph_config["value"] is not workflow.graph_dict - assert [node["id"] for node in captured_graph_config["value"]["nodes"]] == ["node-1", "node-child"] - assert captured_graph_config["value"]["edges"] == [ - {"id": "edge-1", "source": "node-1", "target": "node-child"} - ] - - def test_get_graph_and_variable_pool_for_single_node_run_includes_trace_session_id( - self, monkeypatch: pytest.MonkeyPatch - ): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - graph_config = { - "nodes": [{"id": "node-1", "data": {"type": "start", "version": "1"}}], - "edges": [], - } - workflow = _workflow(graph_config) - workflow.id = "workflow" - captured = {} + def _graph_init(**kwargs): + variable_pool = graph_runtime_state.variable_pool + assert variable_pool.get(["sys", "conversation_id"]) is not None + return SimpleNamespace() - def fake_from_graph_init_context(**kwargs): - captured["run_context"] = kwargs["graph_init_context"].run_context - captured["runtime_state"] = kwargs["runtime_state"] - return SimpleNamespace() + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.Graph.init", + _graph_init, + ) + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.resolve_workflow_node_class", + lambda **_kwargs: _LoopNodeCls, + ) + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.load_into_variable_pool", + lambda **kwargs: None, + ) + monkeypatch.setattr( + "services.workflow.execution.adapters.graph.WorkflowEntry.mapping_user_inputs_to_variable_pool", + lambda **kwargs: None, + ) - class _NodeCls: - @staticmethod - def extract_variable_selector_to_variable_mapping(graph_config, config): - return {} + graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="loop-node", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="loop", + ) - from core.app.apps import workflow_app_runner + assert graph is not None + assert variable_pool.get(["sys", "conversation_id"]).value == "conv-1" - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.DifyNodeFactory.from_graph_init_context", - fake_from_graph_init_context, - ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", lambda **kwargs: SimpleNamespace()) - monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) - monkeypatch.setattr("core.app.apps.workflow_app_runner.load_into_variable_pool", lambda **kwargs: None) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", - lambda **kwargs: None, - ) + def test_handle_graph_run_events_and_pause_notifications(self, monkeypatch: pytest.MonkeyPatch): + published: list[object] = [] - runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="node-1", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="iteration", - user_id="00000000-0000-0000-0000-000000000001", - trace_session_id="session-1", - ) + class _QueueManager: + def publish(self, event, publish_from): + published.append((event, publish_from)) - assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" - assert captured["runtime_state"] is graph_runtime_state - - def test_get_graph_and_variable_pool_preloads_constructor_variables_before_graph_init( - self, monkeypatch: pytest.MonkeyPatch - ): - variable_loader = SimpleNamespace( - load_variables=lambda selectors: ( - [ - StringVariable( - name="conversation_id", - value="conv-1", - selector=["sys", "conversation_id"], - ) - ] - if selectors - else [] - ) - ) - runner = WorkflowBasedAppRunner( - queue_manager=SimpleNamespace(), - variable_loader=variable_loader, - app_id="app", - ) - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) + runner = WorkflowEventPublisher(queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications) + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + workflow_entry = _workflow_entry(runner, graph_runtime_state) - workflow = _workflow( - { - "nodes": [ - {"id": "loop-node", "data": {"type": "loop", "version": "1", "title": "Loop"}}, - { - "id": "llm-child", - "data": { - "type": "llm", - "version": "1", - "loop_id": "loop-node", - "memory": object(), - }, - }, - ], - "edges": [], - } - ) - workflow.id = "workflow" + emails: list[dict] = [] - class _LoopNodeCls: - @staticmethod - def extract_variable_selector_to_variable_mapping(graph_config, config): - return {} + class _Dispatch: + def apply_async(self, *, kwargs, queue): + emails.append({"kwargs": kwargs, "queue": queue}) - def _graph_init(**kwargs): - variable_pool = graph_runtime_state.variable_pool - assert variable_pool.get(["sys", "conversation_id"]) is not None - return SimpleNamespace() + monkeypatch.setattr( + "services.workflow.execution.adapters.events.dispatch_human_input_form_delivery_task", + _Dispatch(), + ) + enriched_reason = HumanInputRequired( + form_id="form", + form_content="content", + node_id="node-1", + node_title="Node", + ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.Graph.init", - _graph_init, - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.resolve_workflow_node_class", - lambda **_kwargs: _LoopNodeCls, - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.load_into_variable_pool", - lambda **kwargs: None, - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", - lambda **kwargs: None, - ) + reason = HitlRequired( + session_id="form", + node_id="node-1", + node_title="Node", + ) - graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="loop-node", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="loop", - ) + runner.publish(workflow_entry, GraphRunStartedEvent()) + runner.publish(workflow_entry, GraphRunSucceededEvent(outputs={"ok": True})) + runner.publish( + workflow_entry, GraphRunPausedEvent(reasons=[reason], outputs={}), pause_reasons=[enriched_reason] + ) - assert graph is not None - assert variable_pool.get(["sys", "conversation_id"]).value == "conv-1" + assert any(isinstance(event, QueueWorkflowStartedEvent) for event, _ in published) + assert any(isinstance(event, QueueWorkflowSucceededEvent) for event, _ in published) + paused_event = next(event for event, _ in published if isinstance(event, QueueWorkflowPausedEvent)) + assert paused_event.paused_nodes == ["node-1"] + assert emails - def test_handle_graph_run_events_and_pause_notifications(self, monkeypatch: pytest.MonkeyPatch): - published: list[object] = [] + def test_handle_graph_aborted_publishes_stopped_terminal(self): + published: list[object] = [] - class _QueueManager: - def publish(self, event, publish_from): - published.append((event, publish_from)) + class _QueueManager: + def publish(self, event, publish_from): + del publish_from + published.append(event) - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - workflow_entry = _workflow_entry(runner, graph_runtime_state) + runner = WorkflowEventPublisher(queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications) + workflow_entry = _workflow_entry(runner) - emails: list[dict] = [] + runner.publish(workflow_entry, GraphRunAbortedEvent(reason="User requested stop", outputs={})) - class _Dispatch: - def apply_async(self, *, kwargs, queue): - emails.append({"kwargs": kwargs, "queue": queue}) + event = published[-1] + assert isinstance(event, QueueStopEvent) + assert event.get_stop_reason() == "User requested stop" - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.dispatch_human_input_form_delivery_task", - _Dispatch(), - ) - enriched_reason = HumanInputRequired( - form_id="form", - form_content="content", - node_id="node-1", - node_title="Node", - ) + def test_handle_node_events_publishes_queue_events(self): + published: list[object] = [] - reason = HitlRequired( - session_id="form", - node_id="node-1", - node_title="Node", - ) + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) - runner.handle_event(workflow_entry, GraphRunStartedEvent()) - runner.handle_event(workflow_entry, GraphRunSucceededEvent(outputs={"ok": True})) - runner.handle_event( - workflow_entry, GraphRunPausedEvent(reasons=[reason], outputs={}), pause_reasons=[enriched_reason] - ) + runner = WorkflowEventPublisher(queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications) + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + workflow_entry = _workflow_entry(runner, graph_runtime_state) + + runner.publish( + workflow_entry, + NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + node_title="Start", + start_at=datetime.now(UTC), + ), + ) + runner.publish( + workflow_entry, + NodeRunStreamChunkEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + selector=["node", "text"], + chunk="hi", + is_final=False, + ), + ) + runner.publish( + workflow_entry, + NodeRunReasoningChunkEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + selector=["node", "reasoning_content"], + chunk="thinking", + is_final=False, + ), + ) + runner.publish( + workflow_entry, + NodeRunAgentLogEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + message_id="msg", + label="label", + node_execution_id="exec", + parent_id=None, + error=None, + status="done", + data={}, + metadata={}, + ), + ) + runner.publish( + workflow_entry, + NodeRunIterationSucceededEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="Iter", + start_at=datetime.now(UTC), + inputs={}, + outputs={"ok": True}, + metadata={}, + steps=1, + ), + ) + runner.publish( + workflow_entry, + NodeRunLoopFailedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="Loop", + start_at=datetime.now(UTC), + inputs={}, + outputs={}, + metadata={}, + steps=1, + error="boom", + ), + ) - assert any(isinstance(event, QueueWorkflowStartedEvent) for event, _ in published) - assert any(isinstance(event, QueueWorkflowSucceededEvent) for event, _ in published) - paused_event = next(event for event, _ in published if isinstance(event, QueueWorkflowPausedEvent)) - assert paused_event.paused_nodes == ["node-1"] - assert emails + assert any(isinstance(event, QueueTextChunkEvent) for event in published) + assert any(isinstance(event, QueueReasoningChunkEvent) for event in published) + assert any(isinstance(event, QueueAgentLogEvent) for event in published) + assert any(isinstance(event, QueueIterationCompletedEvent) for event in published) + assert any(isinstance(event, QueueLoopCompletedEvent) for event in published) - def test_handle_graph_aborted_publishes_stopped_terminal(self): - published: list[object] = [] + def test_handle_human_input_form_filled_event_preserves_submitted_data(self): + published: list[object] = [] - class _QueueManager: - def publish(self, event, publish_from): - del publish_from - published.append(event) + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - workflow_entry = _workflow_entry(runner) + publisher = WorkflowEventPublisher( + queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications + ) + workflow_entry = _workflow_entry() + submitted_data = {"decision": build_segment("approve")} + + publisher.publish( + workflow_entry, + NodeRunHumanInputFormFilledEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.HUMAN_INPUT, + node_title="Human Input", + rendered_content="content", + action_id="approve", + action_text="Approve", + submitted_data=submitted_data, + ), + ) - runner.handle_event(workflow_entry, GraphRunAbortedEvent(reason="User requested stop", outputs={})) + queue_event = published[-1] + assert isinstance(queue_event, QueueHumanInputFormFilledEvent) + assert queue_event.submitted_data == submitted_data - event = published[-1] - assert isinstance(event, QueueStopEvent) - assert event.get_stop_reason() == "User requested stop" + @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) + def test_nested_engine_events_retain_both_container_ancestors(self, ownership: str): + from tests.unit_tests.services.test_workflow_persistence import _make_layer - def test_handle_node_events_publishes_queue_events(self): - published: list[object] = [] + published: list[object] = [] - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) + class _QueueManager: + def publish(self, event, publish_from): + del publish_from + published.append(event) - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - workflow_entry = _workflow_entry(runner, graph_runtime_state) + graph_config = _nested_container_graph(ownership) + graph_config["nodes"].append( + {"id": "deep-end", "data": {"type": "end", "title": "Nested End", "outputs": [], "container_id": "deep"}} + ) + graph_config["edges"].append({"id": "deep-output", "source": "deep-start", "target": "deep-end"}) + layer, _, node_repo, state = _make_layer(graph_data=graph_config) + persisted_starts: dict[str, dict] = {} + + class _StartMetadataLayer(Layer): + @override + def on_event(self, event): + if isinstance(event, NodeRunStartedEvent): + persisted_starts[event.node_id] = dict(node_repo.saved[-1].metadata) + + state.variable_pool.add(["start", "items"], ["item"]) + builder = WorkflowGraphBuilder(app_id="app") + publisher = WorkflowEventPublisher( + queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications + ) + graph = builder.build( + graph_config=graph_config, + graph_runtime_state=state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + human_form_reader = MagicMock() + entry = WorkflowEntry( + graph_config=graph.graph_config, + graph=graph, + call_depth=0, + graph_runtime_state=state, + human_form_reader=human_form_reader, + ) + engine = entry.graph_engine + index = WorkflowRunIndex(()) + layer.set_node_run_indices(index.indices) + engine.add_layer(index) + engine.add_layer(layer) + engine.add_layer(_StartMetadataLayer()) + for event in iter_dify_graph_engine_events(engine, human_form_reader=human_form_reader): + publisher.publish(entry, event) + + expected = { + "outer": (None, None), + "inner": (None, "outer"), + "deep": ("inner", "outer"), + "deep-end": ("inner", "deep"), + } + started = {event.node_id: event for event in published if isinstance(event, QueueNodeStartedEvent)} + persisted = {event.node_id: event for event in node_repo.saved} + for node_id, (iteration_id, loop_id) in expected.items(): + assert (started[node_id].in_iteration_id, started[node_id].in_loop_id) == (iteration_id, loop_id) + for metadata in (persisted_starts[node_id], persisted[node_id].metadata): + assert (metadata.get("iteration_id"), metadata.get("loop_id")) == (iteration_id, loop_id) + for metadata in (persisted_starts["deep-end"], persisted["deep-end"].metadata): + assert metadata["iteration_index"] == 0 + assert metadata["loop_index"] == 0 + for event in published: + if isinstance( + event, + QueueIterationStartEvent | QueueIterationCompletedEvent | QueueLoopStartEvent | QueueLoopCompletedEvent, + ): + assert (event.metadata.get("iteration_id"), event.metadata.get("loop_id")) == expected[event.node_id] - runner.handle_event( - workflow_entry, - NodeRunStartedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - node_title="Start", - start_at=datetime.now(UTC), + def test_handle_nested_answer_success_publishes_text_chunk(self): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) + + runner = WorkflowEventPublisher(queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications) + workflow_entry = _workflow_entry(runner) + + runner.publish( + workflow_entry, + NodeRunSucceededEvent( + id="exec", + node_id="answer-in-iteration", + node_type=BuiltinNodeTypes.ANSWER, + start_at=datetime.now(UTC), + finished_at=datetime.now(UTC), + node_run_result=NodeRunResult( + outputs={"answer": "inside iteration"}, metadata={"iteration_id": "iteration"} ), - ) - runner.handle_event( - workflow_entry, - NodeRunStreamChunkEvent( + container_id="iteration", + ), + ) + + text_chunk = published[0] + assert isinstance(text_chunk, QueueTextChunkEvent) + assert text_chunk.text == "inside iteration" + assert text_chunk.from_variable_selector == ["answer-in-iteration", "answer"] + assert text_chunk.in_iteration_id == "iteration" + assert any(isinstance(event, QueueNodeSucceededEvent) for event in published) + + def test_handle_root_answer_success_does_not_publish_extra_text_chunk(self): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) + + runner = WorkflowEventPublisher(queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications) + workflow_entry = _workflow_entry(runner) + + runner.publish( + workflow_entry, + NodeRunSucceededEvent( + id="exec", + node_id="answer", + node_type=BuiltinNodeTypes.ANSWER, + start_at=datetime.now(UTC), + finished_at=datetime.now(UTC), + node_run_result=NodeRunResult(outputs={"answer": "root answer"}), + ), + ) + + assert not any(isinstance(event, QueueTextChunkEvent) for event in published) + assert len(published) == 1 + assert isinstance(published[0], QueueNodeSucceededEvent) + + @pytest.mark.parametrize( + ("event_factory", "queue_event_cls"), + [ + ( + lambda result, start_at, finished_at: NodeRunSucceededEvent( id="exec", node_id="node", node_type=BuiltinNodeTypes.START, - selector=["node", "text"], - chunk="hi", - is_final=False, + start_at=start_at, + finished_at=finished_at, + node_run_result=result, ), - ) - runner.handle_event( - workflow_entry, - NodeRunReasoningChunkEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - selector=["node", "reasoning_content"], - chunk="thinking", - is_final=False, - ), - ) - runner.handle_event( - workflow_entry, - NodeRunAgentLogEvent( + QueueNodeSucceededEvent, + ), + ( + lambda result, start_at, finished_at: NodeRunFailedEvent( id="exec", node_id="node", node_type=BuiltinNodeTypes.START, - message_id="msg", - label="label", - node_execution_id="exec", - parent_id=None, - error=None, - status="done", - data={}, - metadata={}, + start_at=start_at, + finished_at=finished_at, + error="boom", + node_run_result=result, ), - ) - runner.handle_event( - workflow_entry, - NodeRunIterationSucceededEvent( + QueueNodeFailedEvent, + ), + ( + lambda result, start_at, finished_at: NodeRunExceptionEvent( id="exec", node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="Iter", - start_at=datetime.now(UTC), - inputs={}, - outputs={"ok": True}, - metadata={}, - steps=1, + node_type=BuiltinNodeTypes.START, + start_at=start_at, + finished_at=finished_at, + error="boom", + node_run_result=result, ), - ) - runner.handle_event( - workflow_entry, - NodeRunLoopFailedEvent( + QueueNodeExceptionEvent, + ), + ( + lambda result, start_at, _finished_at: NodeRunRetryEvent( id="exec", node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="Loop", - start_at=datetime.now(UTC), - inputs={}, - outputs={}, - metadata={}, - steps=1, + node_type=BuiltinNodeTypes.START, + node_title="Start", + start_at=start_at, error="boom", + retry_index=1, + node_run_result=result, ), - ) - - assert any(isinstance(event, QueueTextChunkEvent) for event in published) - assert any(isinstance(event, QueueReasoningChunkEvent) for event in published) - assert any(isinstance(event, QueueAgentLogEvent) for event in published) - assert any(isinstance(event, QueueIterationCompletedEvent) for event in published) - assert any(isinstance(event, QueueLoopCompletedEvent) for event in published) - - @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) - def test_nested_engine_events_retain_both_container_ancestors(self, ownership: str): - from tests.unit_tests.services.test_workflow_persistence import _make_layer - - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - del publish_from - published.append(event) - - graph_config = _nested_container_graph(ownership) - graph_config["nodes"].append( - { - "id": "deep-end", - "data": {"type": "end", "title": "Nested End", "outputs": [], "container_id": "deep"}, - } - ) - graph_config["edges"].append({"id": "deep-output", "source": "deep-start", "target": "deep-end"}) - layer, _, node_repo, state = _make_layer(graph_data=graph_config) - persisted_starts: dict[str, dict] = {} - - class _StartMetadataLayer(Layer): - @override - def on_event(self, event): - if isinstance(event, NodeRunStartedEvent): - persisted_starts[event.node_id] = dict(node_repo.saved[-1].metadata) - - state.variable_pool.add(["start", "items"], ["item"]) - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph = runner._init_graph( - graph_config=graph_config, - graph_runtime_state=state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - entry = WorkflowEntry(graph_config=graph.graph_config, graph=graph, call_depth=0, graph_runtime_state=state) - engine = entry.graph_engine - index = WorkflowRunIndex(()) - layer.set_node_run_indices(index.indices) - engine.add_layer(index) - engine.add_layer(layer) - engine.add_layer(_StartMetadataLayer()) - for event in iter_dify_graph_engine_events(engine): - runner.handle_event(entry, event) - - expected = { - "outer": (None, None), - "inner": (None, "outer"), - "deep": ("inner", "outer"), - "deep-end": ("inner", "deep"), - } - started = {event.node_id: event for event in published if isinstance(event, QueueNodeStartedEvent)} - persisted = {event.node_id: event for event in node_repo.saved} - for node_id, (iteration_id, loop_id) in expected.items(): - assert (started[node_id].in_iteration_id, started[node_id].in_loop_id) == (iteration_id, loop_id) - for metadata in (persisted_starts[node_id], persisted[node_id].metadata): - assert (metadata.get("iteration_id"), metadata.get("loop_id")) == (iteration_id, loop_id) - for metadata in (persisted_starts["deep-end"], persisted["deep-end"].metadata): - assert metadata["iteration_index"] == 0 - assert metadata["loop_index"] == 0 - for event in published: - if isinstance( - event, - QueueIterationStartEvent - | QueueIterationCompletedEvent - | QueueLoopStartEvent - | QueueLoopCompletedEvent, - ): - assert (event.metadata.get("iteration_id"), event.metadata.get("loop_id")) == expected[ - event.node_id - ] - - def test_handle_nested_answer_success_publishes_text_chunk(self): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - workflow_entry = _workflow_entry(runner) - - runner.handle_event( - workflow_entry, - NodeRunSucceededEvent( - id="exec", - node_id="answer-in-iteration", - node_type=BuiltinNodeTypes.ANSWER, - start_at=datetime.now(UTC), - finished_at=datetime.now(UTC), - node_run_result=NodeRunResult( - outputs={"answer": "inside iteration"}, metadata={"iteration_id": "iteration"} - ), - container_id="iteration", - ), - ) - - text_chunk = published[0] - assert isinstance(text_chunk, QueueTextChunkEvent) - assert text_chunk.text == "inside iteration" - assert text_chunk.from_variable_selector == ["answer-in-iteration", "answer"] - assert text_chunk.in_iteration_id == "iteration" - assert any(isinstance(event, QueueNodeSucceededEvent) for event in published) - - def test_handle_root_answer_success_does_not_publish_extra_text_chunk(self): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - workflow_entry = _workflow_entry(runner) - - runner.handle_event( - workflow_entry, - NodeRunSucceededEvent( - id="exec", - node_id="answer", - node_type=BuiltinNodeTypes.ANSWER, - start_at=datetime.now(UTC), - finished_at=datetime.now(UTC), - node_run_result=NodeRunResult(outputs={"answer": "root answer"}), - ), - ) - - assert not any(isinstance(event, QueueTextChunkEvent) for event in published) - assert len(published) == 1 - assert isinstance(published[0], QueueNodeSucceededEvent) - - @pytest.mark.parametrize( - ("event_factory", "queue_event_cls"), - [ - ( - lambda result, start_at, finished_at: NodeRunSucceededEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=start_at, - finished_at=finished_at, - node_run_result=result, - ), - QueueNodeSucceededEvent, - ), - ( - lambda result, start_at, finished_at: NodeRunFailedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=start_at, - finished_at=finished_at, - error="boom", - node_run_result=result, - ), - QueueNodeFailedEvent, - ), - ( - lambda result, start_at, finished_at: NodeRunExceptionEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=start_at, - finished_at=finished_at, - error="boom", - node_run_result=result, - ), - QueueNodeExceptionEvent, - ), - ( - lambda result, start_at, _finished_at: NodeRunRetryEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - node_title="Start", - start_at=start_at, - error="boom", - retry_index=1, - node_run_result=result, - ), - QueueNodeRetryEvent, - ), - ], + QueueNodeRetryEvent, + ), + ], + ) + def test_handle_start_node_result_events_project_outputs(self, event_factory, queue_event_cls): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) + + runner = WorkflowEventPublisher(queue_manager=_QueueManager(), notify_pause=enqueue_human_input_notifications) + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + workflow_entry = _workflow_entry(runner, graph_runtime_state) + started_at = datetime.now(UTC) + finished_at = datetime.now(UTC) + result = NodeRunResult( + inputs={"question": "hello"}, + outputs={ + "question": "hello", + "sys.query": "hello", + "env.API_KEY": "secret", + "conversation.session_id": "session-1", + }, ) - def test_handle_start_node_result_events_project_outputs(self, event_factory, queue_event_cls): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - workflow_entry = _workflow_entry(runner, graph_runtime_state) - started_at = datetime.now(UTC) - finished_at = datetime.now(UTC) - result = NodeRunResult( - inputs={"question": "hello"}, - outputs={ - "question": "hello", - "sys.query": "hello", - "env.API_KEY": "secret", - "conversation.session_id": "session-1", - }, - ) - runner.handle_event(workflow_entry, event_factory(result, started_at, finished_at)) + runner.publish(workflow_entry, event_factory(result, started_at, finished_at)) - queue_event = published[-1] - assert isinstance(queue_event, queue_event_cls) - assert queue_event.outputs == {"question": "hello"} + queue_event = published[-1] + assert isinstance(queue_event, queue_event_cls) + assert queue_event.outputs == {"question": "hello"} class _DummyQueueManager: @@ -1049,35 +1090,35 @@ def __init__(self): self.graph_engine = _DummyGraphEngine() -class TestWorkflowAppRunnerNotificationsCases: - @staticmethod - def test_handle_pause_event_enqueues_email_task(monkeypatch: pytest.MonkeyPatch): - queue_manager = _DummyQueueManager() - runner = WorkflowBasedAppRunner(queue_manager=queue_manager, app_id="app-id") - workflow_entry = _DummyWorkflowEntry() +def test_handle_pause_event_enqueues_delivery_task(monkeypatch: pytest.MonkeyPatch): + queue_manager = _DummyQueueManager() + publisher = WorkflowEventPublisher(queue_manager=queue_manager, notify_pause=enqueue_human_input_notifications) + workflow_entry = _DummyWorkflowEntry() - graph_reason = HitlRequired(session_id="form-123", node_id="node-1", node_title="Review") - event = GraphRunPausedEvent(reasons=[graph_reason], outputs={}) + graph_reason = HitlRequired(session_id="form-123", node_id="node-1", node_title="Review") + event = GraphRunPausedEvent(reasons=[graph_reason], outputs={}) - email_task = MagicMock() - enriched_reason = HumanInputRequired( - form_id="form-123", - form_content="content", - inputs=[], - actions=[], - node_id="node-1", - node_title="Review", - ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_form_delivery_task", email_task) + delivery_task = MagicMock() + enriched_reason = HumanInputRequired( + form_id="form-123", + form_content="content", + inputs=[], + actions=[], + node_id="node-1", + node_title="Review", + ) + monkeypatch.setattr( + "services.workflow.execution.adapters.events.dispatch_human_input_form_delivery_task", delivery_task + ) - runner.handle_event(workflow_entry, event, pause_reasons=[enriched_reason]) + publisher.publish(workflow_entry, event, pause_reasons=[enriched_reason]) - email_task.apply_async.assert_called_once() - kwargs = email_task.apply_async.call_args.kwargs["kwargs"] - assert kwargs["form_id"] == "form-123" - assert kwargs["node_title"] == "Review" + delivery_task.apply_async.assert_called_once() + kwargs = delivery_task.apply_async.call_args.kwargs["kwargs"] + assert kwargs["form_id"] == "form-123" + assert kwargs["node_title"] == "Review" - assert any(isinstance(evt, QueueWorkflowPausedEvent) for evt in queue_manager.published) + assert any(isinstance(evt, QueueWorkflowPausedEvent) for evt in queue_manager.published) def _make_graph_state(): @@ -1094,241 +1135,257 @@ def _make_graph_state(): ) -class TestWorkflowAppRunnerSingleNodeCases: - @staticmethod - @pytest.fixture - def queue_manager() -> WorkflowAppQueueManager: - return WorkflowAppQueueManager( - task_id="task-id", user_id="user", invoke_from=InvokeFrom.SERVICE_API, app_mode=AppMode.WORKFLOW - ) +@pytest.fixture +def queue_manager() -> WorkflowAppQueueManager: + return WorkflowAppQueueManager( + task_id="task-id", user_id="user", invoke_from=InvokeFrom.SERVICE_API, app_mode=AppMode.WORKFLOW + ) - @staticmethod - @pytest.mark.parametrize( - ("single_iteration_run", "single_loop_run"), - [ - (WorkflowAppGenerateEntity.SingleIterationRunEntity(node_id="iter", inputs={}), None), - (None, WorkflowAppGenerateEntity.SingleLoopRunEntity(node_id="loop", inputs={})), - ], + +@pytest.mark.parametrize( + ("single_iteration_run", "single_loop_run"), + [ + (WorkflowAppGenerateEntity.SingleIterationRunEntity(node_id="iter", inputs={}), None), + (None, WorkflowAppGenerateEntity.SingleLoopRunEntity(node_id="loop", inputs={})), + ], +) +def test_run_uses_single_node_execution_branch( + single_iteration_run: WorkflowAppGenerateEntity.SingleIterationRunEntity | None, + single_loop_run: WorkflowAppGenerateEntity.SingleLoopRunEntity | None, + queue_manager: WorkflowAppQueueManager, + *, + workflow_runtime: WorkflowRuntime, +) -> None: + app_generate_entity = WorkflowAppGenerateEntity( + user_from=UserFrom.END_USER, + app_config=WorkflowUIBasedAppConfig( + app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW + ), + inputs={}, + files=[], + user_id="user", + invoke_from=InvokeFrom.SERVICE_API, + workflow_execution_id="execution-id", + task_id="task-id", + stream=True, + extras={"trace_session_id": "session-1"}, + single_iteration_run=single_iteration_run, + single_loop_run=single_loop_run, ) - def test_run_uses_single_node_execution_branch( - single_iteration_run: WorkflowAppGenerateEntity.SingleIterationRunEntity | None, - single_loop_run: WorkflowAppGenerateEntity.SingleLoopRunEntity | None, - queue_manager: WorkflowAppQueueManager, - ) -> None: - app_generate_entity = WorkflowAppGenerateEntity( - user_from=UserFrom.END_USER, - app_config=WorkflowUIBasedAppConfig( - app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW - ), - inputs={}, - files=[], - user_id="user", - invoke_from=InvokeFrom.SERVICE_API, - workflow_execution_id="execution-id", - task_id="task-id", - stream=True, - extras={"trace_session_id": "session-1"}, - single_iteration_run=single_iteration_run, - single_loop_run=single_loop_run, - ) - - workflow = Workflow( - tenant_id="tenant", - app_id="app", - id="workflow", - type=ModelWorkflowType.WORKFLOW, - version="v1", - graph=json.dumps({"nodes": [], "edges": []}), - ) - workflow.environment_variables = [] - runner = WorkflowAppRunner( - application_generate_entity=app_generate_entity, - queue_manager=queue_manager, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id="system-user", - ) - - graph, variable_pool, graph_runtime_state = _make_graph_state() - mock_workflow_entry = MagicMock() - mock_workflow_entry.graph_engine = MagicMock() - mock_workflow_entry.graph_engine.layer = MagicMock() - mock_workflow_entry.run.return_value = iter([]) - - with ( - patch("core.app.apps.workflow.app_runner.RedisChannel"), - patch("core.app.apps.workflow.app_runner.redis_client"), - patch("core.app.apps.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry) as entry_class, - patch.object( - runner, - "_prepare_single_node_execution", - return_value=( - graph, - variable_pool, - graph_runtime_state, - ), - ) as prepare_single, - patch.object(runner, "_init_graph") as init_graph, - ): - runner.prepare() + workflow = Workflow( + tenant_id="tenant", + app_id="app", + id="workflow", + type=ModelWorkflowType.WORKFLOW, + version="v1", + graph=json.dumps({"nodes": [], "edges": []}), + ) + workflow.environment_variables = [] + snapshot = workflow_snapshot(workflow) + + runner = WorkflowAppRunner( + application_generate_entity=app_generate_entity, + queue_manager=queue_manager, + variable_loader=MagicMock(), + workflow=snapshot, + system_user_id="system-user", + runtime=workflow_runtime, + ) - prepare_single.assert_called_once_with( - workflow=workflow, - single_iteration_run=single_iteration_run, - single_loop_run=single_loop_run, - user_id="user", - app_type=CreditUsageAppType.WORKFLOW, - trace_session_id="session-1", - ) - init_graph.assert_not_called() + graph, variable_pool, graph_runtime_state = _make_graph_state() + mock_workflow_entry = MagicMock() + mock_workflow_entry.graph_engine = MagicMock() + mock_workflow_entry.graph_engine.layer = MagicMock() + mock_workflow_entry.run.return_value = iter([]) + + with ( + patch("services.workflow.execution.adapters.workflow.app_runner.RedisChannel"), + patch("services.workflow.execution.adapters.workflow.app_runner.redis_client"), + patch( + "services.workflow.execution.adapters.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry + ) as entry_class, + patch.object( + runner._graphs, + "build_single_node", + return_value=( + graph, + variable_pool, + graph_runtime_state, + ), + ) as prepare_single, + patch.object(runner._graphs, "build") as init_graph, + ): + runner.prepare() + + prepare_single.assert_called_once_with( + workflow=snapshot, + single_iteration_run=single_iteration_run, + single_loop_run=single_loop_run, + user_id="user", + app_type=CreditUsageAppType.WORKFLOW, + trace_session_id="session-1", + ) + init_graph.assert_not_called() - assert entry_class.call_args.kwargs["graph_runtime_state"] is graph_runtime_state + assert entry_class.call_args.kwargs["graph_runtime_state"] is graph_runtime_state - @staticmethod - def test_single_node_run_validates_target_node_config(queue_manager: WorkflowAppQueueManager) -> None: - runner = WorkflowBasedAppRunner( - queue_manager=queue_manager, - variable_loader=MagicMock(), - app_id="app", - ) - workflow = Workflow( - id="workflow", - tenant_id="tenant", - graph=json.dumps( - { - "nodes": [ - { - "id": "loop-node", - "data": {}, - } - ], - "edges": [], - } - ), - ) +def test_single_node_run_validates_target_node_config() -> None: + runner = WorkflowGraphBuilder( + variable_loader=MagicMock(), + app_id="app", + ) - _, _, graph_runtime_state = _make_graph_state() + workflow = Workflow( + id="workflow", + tenant_id="tenant", + graph=json.dumps( + { + "nodes": [ + { + "id": "loop-node", + "data": {}, + } + ], + "edges": [], + } + ), + ) - with pytest.raises(ValidationError, match="nodes.0.data.type"): - runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="loop-node", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="loop", - user_id="00000000-0000-0000-0000-000000000001", - ) + _, _, graph_runtime_state = _make_graph_state() - @staticmethod - def test_single_node_run_rejects_invalid_loop_count(queue_manager: WorkflowAppQueueManager) -> None: - runner = WorkflowBasedAppRunner( - queue_manager=queue_manager, - app_id="app", + with pytest.raises(ValidationError, match="nodes.0.data.type"): + runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow_snapshot(workflow), + node_id="loop-node", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="loop", + user_id="00000000-0000-0000-0000-000000000001", ) - workflow = Workflow( - id="workflow", - tenant_id="tenant", - graph=json.dumps( - { - "nodes": [ - { - "id": "loop-node", - "data": { - "type": "loop", - "title": "Loop", - "loop_count": 0, - "start_node_id": "loop-start", - "break_conditions": [], - "logical_operator": "and", - }, - }, - { - "id": "loop-start", - "data": {"type": "loop-start", "title": "Loop start", "container_id": "loop-node"}, + +def test_single_node_run_rejects_invalid_loop_count() -> None: + runner = WorkflowGraphBuilder(app_id="app") + + workflow = Workflow( + id="workflow", + tenant_id="tenant", + graph=json.dumps( + { + "nodes": [ + { + "id": "loop-node", + "data": { + "type": "loop", + "title": "Loop", + "loop_count": 0, + "start_node_id": "loop-start", + "break_conditions": [], + "logical_operator": "and", }, - ], - "edges": [], - } - ), + }, + { + "id": "loop-start", + "data": {"type": "loop-start", "title": "Loop start", "container_id": "loop-node"}, + }, + ], + "edges": [], + } + ), + ) + + _, _, graph_runtime_state = _make_graph_state() + with pytest.raises(ValidationError, match="loop_count"): + runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow_snapshot(workflow), + node_id="loop-node", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="loop", + user_id="00000000-0000-0000-0000-000000000001", ) - _, _, graph_runtime_state = _make_graph_state() - with pytest.raises(ValidationError, match="loop_count"): - runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="loop-node", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="loop", - user_id="00000000-0000-0000-0000-000000000001", - ) - @staticmethod - def test_run_adds_inputs_with_snippet_compatible_start_aliases(queue_manager: WorkflowAppQueueManager) -> None: - app_generate_entity = WorkflowAppGenerateEntity( - user_from=UserFrom.END_USER, - app_config=WorkflowUIBasedAppConfig( - app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW - ), - inputs={"question": "hello"}, - files=[], - user_id="user", - invoke_from=InvokeFrom.SERVICE_API, - workflow_execution_id="execution-id", - task_id="task-id", - stream=True, - call_depth=4, - ) +def test_run_adds_inputs_with_snippet_compatible_start_aliases( + queue_manager: WorkflowAppQueueManager, *, workflow_runtime: WorkflowRuntime +) -> None: + app_generate_entity = WorkflowAppGenerateEntity( + user_from=UserFrom.END_USER, + app_config=WorkflowUIBasedAppConfig( + app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW + ), + inputs={"question": "hello"}, + files=[], + user_id="user", + invoke_from=InvokeFrom.SERVICE_API, + workflow_execution_id="execution-id", + task_id="task-id", + stream=True, + call_depth=4, + ) - workflow = Workflow( - tenant_id="tenant", - app_id="app", - id="workflow", - type=ModelWorkflowType.WORKFLOW, - version="v1", - graph=json.dumps({"nodes": [], "edges": []}), - kind=WorkflowKind.SNIPPET, - ) - workflow.environment_variables = [] - - runner = WorkflowAppRunner( - application_generate_entity=app_generate_entity, - queue_manager=queue_manager, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id="system-user", - ) - - mock_workflow_entry = MagicMock() - mock_workflow_entry.graph_engine = MagicMock() - mock_workflow_entry.graph_engine.layer = MagicMock() - mock_workflow_entry.run.return_value = iter([]) - - with ( - patch("core.app.apps.workflow.app_runner.RedisChannel"), - patch("core.app.apps.workflow.app_runner.redis_client"), - patch("core.app.apps.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry), - patch("core.app.apps.workflow.app_runner.build_system_variables", return_value={}), - patch("core.app.apps.workflow.app_runner.build_bootstrap_variables", return_value=[]), - patch("core.app.apps.workflow.app_runner.add_variables_to_pool"), - patch("core.app.apps.workflow.app_runner.get_default_root_node_id", return_value="root-node"), - patch( - "core.app.apps.workflow.app_runner.get_compatible_start_aliases", return_value=("legacy-start",) - ) as aliases, - patch("core.app.apps.workflow.app_runner.add_node_inputs_to_pool") as add_inputs, - patch.object(runner, "_init_graph", return_value=MagicMock()) as init_graph, - ): - runner.prepare() - - aliases.assert_called_once_with(workflow_kind="snippet", root_node_id="root-node") - add_inputs.assert_called_once() - assert add_inputs.call_args.kwargs["node_id"] == "root-node" - assert add_inputs.call_args.kwargs["inputs"] == {"question": "hello"} - assert add_inputs.call_args.kwargs["aliases"] == ("legacy-start",) - assert init_graph.call_args.kwargs["call_depth"] == 4 + workflow = Workflow( + tenant_id="tenant", + app_id="app", + id="workflow", + type=ModelWorkflowType.WORKFLOW, + version="v1", + graph=json.dumps({"nodes": [], "edges": []}), + kind=WorkflowKind.SNIPPET, + ) + workflow.environment_variables = [] + snapshot = workflow_snapshot(workflow) + + runner = WorkflowAppRunner( + application_generate_entity=app_generate_entity, + queue_manager=queue_manager, + variable_loader=MagicMock(), + workflow=snapshot, + system_user_id="system-user", + runtime=workflow_runtime, + ) + + mock_workflow_entry = MagicMock() + mock_workflow_entry.graph_engine = MagicMock() + mock_workflow_entry.graph_engine.layer = MagicMock() + mock_workflow_entry.run.return_value = iter([]) + + with ( + patch("services.workflow.execution.adapters.workflow.app_runner.RedisChannel"), + patch("services.workflow.execution.adapters.workflow.app_runner.redis_client"), + patch( + "services.workflow.execution.adapters.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry + ), + patch("services.workflow.execution.adapters.workflow.app_runner.build_system_variables", return_value={}), + patch("services.workflow.execution.adapters.workflow.app_runner.build_bootstrap_variables", return_value=[]), + patch("services.workflow.execution.adapters.workflow.app_runner.add_variables_to_pool"), + patch( + "services.workflow.execution.adapters.workflow.app_runner.get_default_root_node_id", + return_value="root-node", + ), + patch( + "services.workflow.execution.adapters.workflow.app_runner.get_compatible_start_aliases", + return_value=("legacy-start",), + ) as aliases, + patch("services.workflow.execution.adapters.workflow.app_runner.add_node_inputs_to_pool") as add_inputs, + patch.object(runner._graphs, "build", return_value=MagicMock()) as init_graph, + ): + runner.prepare() + + aliases.assert_called_once_with(workflow_kind="snippet", root_node_id="root-node") + add_inputs.assert_called_once() + assert add_inputs.call_args.kwargs["node_id"] == "root-node" + assert add_inputs.call_args.kwargs["inputs"] == {"question": "hello"} + assert add_inputs.call_args.kwargs["aliases"] == ("legacy-start",) + assert init_graph.call_args.kwargs["call_depth"] == 4 + + +@pytest.fixture +def workflow_queue(mocker: MockerFixture) -> tuple[WorkflowAppQueueManager, MagicMock]: + queue = WorkflowAppQueueManager("task", "user", InvokeFrom.DEBUGGER, AppMode.WORKFLOW) + return queue, mocker.spy(queue, "publish") def _save_form(session: Session, *, status: HumanInputFormStatus) -> HumanInputForm: @@ -1361,13 +1418,67 @@ def _save_form(session: Session, *, status: HumanInputFormStatus) -> HumanInputF return form +def _generate_entity() -> WorkflowAppGenerateEntity: + return WorkflowAppGenerateEntity( + task_id="task", + app_config=WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + workflow_id="workflow", + ), + inputs={}, + files=[], + user_id="user", + stream=True, + invoke_from=InvokeFrom.DEBUGGER, + workflow_execution_id="run-1", + ) + + +def _advanced_chat_generate_entity() -> AdvancedChatAppGenerateEntity: + return AdvancedChatAppGenerateEntity( + task_id="task-1", + app_config=WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow", + ), + inputs={}, + query="query", + files=[], + user_id="user", + stream=True, + invoke_from=InvokeFrom.DEBUGGER, + ) + + +def _build_runner(queue: WorkflowAppQueueManager, runtime: WorkflowExecutionDependencies) -> WorkflowAppRunner: + workflow = make_workflow( + workflow_id="workflow", + tenant_id="tenant", + app_id="app", + graph={"nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], "edges": []}, + created_by="user", + ) + return WorkflowAppRunner( + runtime=runtime, + application_generate_entity=_generate_entity(), + queue_manager=queue, + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=workflow_snapshot(workflow), + system_user_id="user", + ) + + def _make_paused_workflow( - runner: WorkflowBasedAppRunner, forms: list[HumanInputForm], *, human_input_form: HumanInputForm | None = None + runner: WorkflowAppRunner, forms: list[HumanInputForm], *, human_input_form: HumanInputForm | None = None ) -> WorkflowEntry: state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=perf_counter()) state.variable_pool.add(("sys", "workflow_run_id"), "run-1") state.variable_pool.add(("sys", "app_id"), "app") - graph = runner._init_graph( + graph = runner._graphs.build( graph_config={ "nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], "edges": [], @@ -1378,6 +1489,7 @@ def _make_paused_workflow( user_from=UserFrom.END_USER, invoke_from=InvokeFrom.WEB_APP, call_depth=0, + workflow_id="workflow", ) if human_input_form is not None: node_data = HumanInputNodeData( @@ -1387,8 +1499,12 @@ def _make_paused_workflow( user_actions=[UserActionConfig(id="approve", title="Approve")], ) callback = DifyHITLCallback( - form_repository=HumanInputFormRepositoryImpl( - tenant_id="tenant", app_id="app", workflow_execution_id="run-1" + form_repository=runner._runtime.human_forms( + tenant_id="tenant", + app_id="app", + workflow_execution_id="run-1", + invoke_source=InvokeFrom.DEBUGGER.value, + submission_actor_id="user", ), node_data=node_data, execution_id_getter=lambda: human_input_form.id, @@ -1415,46 +1531,28 @@ def _make_paused_workflow( state.graph_execution.start() for form in forms: state.graph_execution.pause(HitlRequired(session_id=form.id, node_id=form.node_id, node_title="Approval")) - state.variable_pool.add(human_input_container_selector(form.id), "outer-tool") state.defer_ready_task(StartTask(frame_id=ROOT_FRAME_ID, node_id=graph.root_node.id)) return WorkflowEntry( graph=graph, graph_config=graph.graph_config or {}, call_depth=0, graph_runtime_state=state, + human_form_reader=runner.human_form_reader, ) -def _resume_events(runner: WorkflowBasedAppRunner, entry: WorkflowEntry) -> Generator[EngineEvent]: +def _resume_events(runner: WorkflowAppRunner, entry: WorkflowEntry) -> Generator[EngineEvent]: """Resume through orchestration with actual repositories and the real Engine.""" - entity = WorkflowAppGenerateEntity( - user_from=UserFrom.ACCOUNT, - task_id="task", - app_config=WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - workflow_id="workflow", - ), - inputs={}, - files=[], - user_id="user", - stream=True, - invoke_from=InvokeFrom.DEBUGGER, - workflow_execution_id="run-1", - ) + entity = runner.application_generate_entity user = Account(name="Tester", email="tester@example.com") user.id = "user" - factory = session_factory.get_session_maker() - runs = SQLAlchemyWorkflowExecutionRepository( - factory, + runs = runner._runtime.execution_writer( tenant_id="tenant", app_id="app", user=user, triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, ) - nodes = SQLAlchemyWorkflowNodeExecutionRepository( - factory, + nodes = runner._runtime.node_writer( tenant_id="tenant", app_id="app", user=user, @@ -1470,272 +1568,293 @@ def _resume_events(runner: WorkflowBasedAppRunner, entry: WorkflowEntry) -> Gene ).iter_events() -class TestWorkflowHumanInputCompletionCases: - @staticmethod - @pytest.fixture - def workflow_queue(mocker: MockerFixture) -> tuple[WorkflowAppQueueManager, MagicMock]: - queue = WorkflowAppQueueManager("task", "user", InvokeFrom.DEBUGGER, AppMode.WORKFLOW) - return queue, mocker.spy(queue, "publish") - - @staticmethod - @pytest.mark.parametrize("status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) - @pytest.mark.parametrize("runs_human_input_node", [False, True], ids=["hidden-form", "direct-completion"]) - def test_resume_publishes_only_the_completed_form_among_repeated_node_invocations( - sqlite_session: Session, - workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], - status: HumanInputFormStatus, - runs_human_input_node: bool, - ) -> None: - selected = _save_form(sqlite_session, status=status) - waiting = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) - expired = _save_form(sqlite_session, status=HumanInputFormStatus.EXPIRED) - _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow( - runner, [selected, waiting, selected, expired], human_input_form=selected if runs_human_input_node else None - ) - original_reasons = tuple(entry.graph_engine.runtime_state.graph_execution.pause_reasons) - - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - published = [call.args[0] for call in publish_spy.call_args_list] - assert isinstance(published[0], QueueWorkflowStartedEvent) - completions = [ - item - for item in published - if isinstance(item, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) - ] - assert len(completions) == 1 - completion = completions[0] - assert completion.form_id == selected.id - assert completion.node_id == "outer-tool" - if status == HumanInputFormStatus.SUBMITTED: - assert isinstance(completion, QueueHumanInputFormFilledEvent) - assert (completion.action_id, completion.action_text, completion.rendered_content) == ( - "approve", - "Approve", - "Decision: approved", - ) - assert completion.submitted_data is not None - assert completion.submitted_data["answer"].to_object() == "approved" - else: - assert isinstance(completion, QueueHumanInputFormTimeoutEvent) - assert len(original_reasons) == 4 - assert all( - isinstance(reason, HitlRequired) and reason.node_id == "same-human-node" for reason in original_reasons - ) - assert entry.graph_engine.runtime_state.graph_execution.pause_reasons == [] - - @staticmethod - @pytest.mark.parametrize("owner_field", ["tenant_id", "app_id", "workflow_run_id"]) - def test_resume_rejects_form_outside_the_trusted_execution_owner( - sqlite_session: Session, workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], owner_field: str - ) -> None: - form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - setattr(form, owner_field, "another-owner") - sqlite_session.commit() - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [form]) - - events = list(_resume_events(runner, entry)) - assert isinstance(events[-1], GraphRunFailedEvent) - assert "does not belong" in events[-1].error - - assert not any(isinstance(call.args[0], QueueHumanInputFormFilledEvent) for call in publish_spy.call_args_list) - - @staticmethod - def test_resume_refreshes_expired_file_and_file_list_urls_before_form_completion( - sqlite_session: Session, - workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], - monkeypatch: pytest.MonkeyPatch, - mocker: MockerFixture, - ) -> None: - mocker.patch.object( - type(workflow_response_converter.db), - "engine", - new_callable=PropertyMock, - return_value=sqlite_session.get_bind(), - ) - submitted_at = 1_700_000_000 - monkeypatch.setattr("core.app.workflow.file_runtime.time.time", lambda: submitted_at) - factory = DifyFileReferenceFactory(build_test_run_context()) - saved_files = [] - for name in ("first.txt", "second.txt"): - upload = UploadFile( - tenant_id="tenant", - storage_type=StorageType.LOCAL, - key=name, - name=name, - size=12, - extension="txt", - mime_type="text/plain", - created_by="user", - created_by_role=CreatorUserRole.END_USER, - created_at=datetime.now(UTC), - used=False, - ) - sqlite_session.add(upload) - sqlite_session.commit() - saved_files.append( - factory.build_from_mapping( - mapping={"transfer_method": "local_file", "upload_file_id": upload.id} - ).to_dict() - ) +@pytest.mark.parametrize("status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) +@pytest.mark.parametrize("runs_human_input_node", [False, True], ids=["hidden-form", "direct-completion"]) +def test_resume_publishes_only_the_completed_form_among_repeated_node_invocations( + sqlite_session: Session, + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + workflow_runtime: WorkflowExecutionDependencies, + status: HumanInputFormStatus, + runs_human_input_node: bool, +) -> None: + selected = _save_form(sqlite_session, status=status) + waiting = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) + expired = _save_form(sqlite_session, status=HumanInputFormStatus.EXPIRED) + _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + queue, publish_spy = workflow_queue + runner = _build_runner(queue, workflow_runtime) + entry = _make_paused_workflow( + runner, [selected, waiting, selected, expired], human_input_form=selected if runs_human_input_node else None + ) + original_reasons = tuple(entry.graph_engine.runtime_state.graph_execution.pause_reasons) - def url_is_valid(file: Mapping[str, object]) -> bool: - url, file_id = file["url"], file["related_id"] - assert isinstance(url, str) - assert isinstance(file_id, str) - query = parse_qs(urlparse(url).query) - return verify_file_signature( - upload_file_id=file_id, - timestamp=query["timestamp"][0], - nonce=query["nonce"][0], - sign=query["sign"][0], - ) + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) - assert all(url_is_valid(file) for file in saved_files) - form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - definition = FormDefinition.model_validate_json(form.form_definition) - definition.inputs.extend( - [ - FileInputConfig(output_variable_name="attachment"), - FileListInputConfig(output_variable_name="attachments"), - ] + published = [call.args[0] for call in publish_spy.call_args_list] + assert isinstance(published[0], QueueWorkflowStartedEvent) + completions = [ + item for item in published if isinstance(item, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) + ] + assert len(completions) == 1 + completion = completions[0] + assert completion.form_id == selected.id + assert completion.node_id == "same-human-node" + if status == HumanInputFormStatus.SUBMITTED: + assert isinstance(completion, QueueHumanInputFormFilledEvent) + assert (completion.action_id, completion.action_text, completion.rendered_content) == ( + "approve", + "Approve", + "Decision: approved", ) - form.form_definition = definition.model_dump_json() - form.submitted_data = json.dumps( - {"answer": "approved", "attachment": saved_files[0], "attachments": saved_files} + assert completion.submitted_data is not None + assert completion.submitted_data["answer"].to_object() == "approved" + else: + assert isinstance(completion, QueueHumanInputFormTimeoutEvent) + assert len(original_reasons) == 4 + assert all(isinstance(reason, HitlRequired) and reason.node_id == "same-human-node" for reason in original_reasons) + assert entry.graph_engine.runtime_state.graph_execution.pause_reasons == [] + + +@pytest.mark.parametrize("owner_field", ["tenant_id", "app_id", "workflow_run_id"]) +def test_resume_rejects_form_outside_the_trusted_execution_owner( + sqlite_session: Session, + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + workflow_runtime: WorkflowExecutionDependencies, + owner_field: str, +) -> None: + form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + setattr(form, owner_field, "another-owner") + sqlite_session.commit() + queue, publish_spy = workflow_queue + runner = _build_runner(queue, workflow_runtime) + entry = _make_paused_workflow(runner, [form]) + + events = list(_resume_events(runner, entry)) + assert isinstance(events[-1], GraphRunFailedEvent) + assert "does not belong" in events[-1].error + + assert not any(isinstance(call.args[0], QueueHumanInputFormFilledEvent) for call in publish_spy.call_args_list) + + +def test_resume_refreshes_expired_file_and_file_list_urls_before_form_completion( + sqlite_session: Session, + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + workflow_runtime: WorkflowExecutionDependencies, + workflow_contexts: ExecutionContexts, + tool_providers: ToolProviders, + monkeypatch: pytest.MonkeyPatch, +) -> None: + submitted_at = 1_700_000_000 + monkeypatch.setattr("core.app.workflow.file_runtime.time.time", lambda: submitted_at) + factory = DifyFileReferenceFactory(build_test_run_context()) + saved_files = [] + for name in ("first.txt", "second.txt"): + upload = UploadFile( + tenant_id="tenant", + storage_type=StorageType.LOCAL, + key=name, + name=name, + size=12, + extension="txt", + mime_type="text/plain", + created_by="user", + created_by_role=CreatorUserRole.END_USER, + created_at=datetime.now(UTC), + used=False, ) + sqlite_session.add(upload) sqlite_session.commit() + saved_files.append( + factory.build_from_mapping(mapping={"transfer_method": "local_file", "upload_file_id": upload.id}).to_dict() + ) - monkeypatch.setattr( - "core.app.workflow.file_runtime.time.time", lambda: submitted_at + dify_config.FILES_ACCESS_TIMEOUT + 1 - ) - assert not any(url_is_valid(file) for file in saved_files) - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [form]) - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - completions = [ - call.args[0] - for call in publish_spy.call_args_list - if isinstance(call.args[0], QueueHumanInputFormFilledEvent) - ] - assert len(completions) == 1 - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task", workflow_run_id="run-1", workflow_id="workflow", reason=WorkflowStartReason.RESUMPTION - ) - response = converter.human_input_form_filled_to_stream_response(event=completions[0], task_id="task") - data = response.data.submitted_data - assert data is not None - assert data["answer"] == "approved" - assert response.data.rendered_content == "Decision: approved" - restored_files = [data["attachment"], *data["attachments"]] - for restored, saved in zip(restored_files, [saved_files[0], *saved_files], strict=True): - assert (restored["related_id"], restored["filename"]) == (saved["related_id"], saved["filename"]) - assert url_is_valid(restored) - - @staticmethod - @pytest.mark.parametrize("late_status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) - def test_form_completed_during_resume_is_published_once_before_the_terminal_event( - sqlite_session: Session, - workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], - late_status: HumanInputFormStatus, - ) -> None: - first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - second = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [first, second]) - - def complete_second_form_during_execution(event: AppQueueEvent, publish_from: PublishFrom) -> None: - WorkflowAppQueueManager.publish(queue, event, publish_from) - # The second submission arrives after the initial completion query, - # while this real Engine attempt is already processing node events. - if not isinstance(event, QueueNodeStartedEvent): - return - repository = HumanInputFormSubmissionRepository() - if late_status == HumanInputFormStatus.SUBMITTED: - repository.mark_submitted( - form_id=second.id, - recipient_id=None, - selected_action_id="approve", - form_data={"answer": "approved"}, - submission_user_id=None, - submission_end_user_id="user", - ) - else: - repository.mark_timeout(form_id=second.id, timeout_status=late_status) - - publish_spy.side_effect = complete_second_form_during_execution - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - published = [call.args[0] for call in publish_spy.call_args_list] - completions = [ - event - for event in published - if isinstance(event, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) - ] - assert [event.form_id for event in completions] == [first.id, second.id] - assert isinstance(published[0], QueueWorkflowStartedEvent) - assert isinstance(published[-1], QueueWorkflowSucceededEvent) - assert published.index(completions[-1]) < len(published) - 1 - - @staticmethod - def test_waiting_form_expiring_during_resume_publishes_timeout_before_success( - sqlite_session: Session, workflow_queue: tuple[WorkflowAppQueueManager, MagicMock] - ) -> None: - form = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) - form.expiration_time = datetime.now(UTC) - timedelta(seconds=1) - sqlite_session.commit() - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [form], human_input_form=form) - - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - published = [call.args[0] for call in publish_spy.call_args_list] - assert isinstance(published[-1], QueueWorkflowSucceededEvent) - timeouts = [event for event in published if isinstance(event, QueueHumanInputFormTimeoutEvent)] - assert [event.form_id for event in timeouts] == [form.id] - assert timeouts[0].node_id == "outer-tool" - assert published.index(timeouts[0]) < len(published) - 1 - sqlite_session.refresh(form) - assert form.status == HumanInputFormStatus.TIMEOUT - - @staticmethod - def test_chat_history_keeps_repeated_forms_by_form_id(sqlite_session: Session) -> None: - pipeline = _build_pipeline() - pipeline._workflow_tenant_id = "tenant" - pipeline._application_generate_entity = _build_generate_entity() - pipeline._workflow_response_converter = _build_converter() - pipeline._workflow_response_converter.workflow_start_to_stream_response( - task_id="task-1", workflow_run_id="run-1", workflow_id="workflow-1", reason=WorkflowStartReason.INITIAL - ) - for form_id in ("first-form", "second-form", "first-form"): - event = QueueHumanInputFormFilledEvent( - form_id=form_id, - node_id="same-human-node", - node_type=BuiltinNodeTypes.HUMAN_INPUT, - node_title="Approval", - rendered_content="Decision: approved", - action_id="approve", - action_text="Approve", + def url_is_valid(file: Mapping[str, object]) -> bool: + url, file_id = file["url"], file["related_id"] + assert isinstance(url, str) + assert isinstance(file_id, str) + query = parse_qs(urlparse(url).query) + return verify_file_signature( + upload_file_id=file_id, + timestamp=query["timestamp"][0], + nonce=query["nonce"][0], + sign=query["sign"][0], + ) + + assert all(url_is_valid(file) for file in saved_files) + form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + definition = FormDefinition.model_validate_json(form.form_definition) + definition.inputs.extend( + [FileInputConfig(output_variable_name="attachment"), FileListInputConfig(output_variable_name="attachments")] + ) + form.form_definition = definition.model_dump_json() + form.submitted_data = json.dumps({"answer": "approved", "attachment": saved_files[0], "attachments": saved_files}) + sqlite_session.commit() + + monkeypatch.setattr( + "core.app.workflow.file_runtime.time.time", lambda: submitted_at + dify_config.FILES_ACCESS_TIMEOUT + 1 + ) + assert not any(url_is_valid(file) for file in saved_files) + queue, publish_spy = workflow_queue + runner = _build_runner(queue, workflow_runtime) + entry = _make_paused_workflow(runner, [form]) + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + + completions = [ + call.args[0] for call in publish_spy.call_args_list if isinstance(call.args[0], QueueHumanInputFormFilledEvent) + ] + assert len(completions) == 1 + converter = _build_converter(workflow_contexts=workflow_contexts, tool_providers=tool_providers) + converter.workflow_start_to_stream_response( + task_id="task", workflow_run_id="run-1", workflow_id="workflow", reason=WorkflowStartReason.RESUMPTION + ) + response = converter.human_input_form_filled_to_stream_response(event=completions[0], task_id="task") + data = response.data.submitted_data + assert data is not None + assert data["answer"] == "approved" + assert response.data.rendered_content == "Decision: approved" + restored_files = [data["attachment"], *data["attachments"]] + for restored, saved in zip(restored_files, [saved_files[0], *saved_files], strict=True): + assert (restored["related_id"], restored["filename"]) == (saved["related_id"], saved["filename"]) + assert url_is_valid(restored) + + +@pytest.mark.parametrize("late_status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) +def test_form_completed_during_resume_is_published_once_before_the_terminal_event( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + workflow_runtime: WorkflowExecutionDependencies, + late_status: HumanInputFormStatus, +) -> None: + first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + second = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) + queue, publish_spy = workflow_queue + runner = _build_runner(queue, workflow_runtime) + entry = _make_paused_workflow(runner, [first, second]) + + def complete_second_form_during_execution(event: AppQueueEvent, publish_from: PublishFrom) -> None: + WorkflowAppQueueManager.publish(queue, event, publish_from) + # The second submission arrives after the initial completion query, + # while this real Engine attempt is already processing node events. + if not isinstance(event, QueueNodeStartedEvent): + return + repository = HumanInputFormSubmissionRepository(sessions=sqlite_session_factory) + if late_status == HumanInputFormStatus.SUBMITTED: + repository.mark_submitted( + form_id=second.id, + recipient_id=None, + selected_action_id="approve", + form_data={"answer": "approved"}, + submission_user_id=None, + submission_end_user_id="user", ) - list(pipeline._handle_human_input_form_filled_event(event)) + else: + repository.mark_timeout(form_id=second.id, timeout_status=late_status) - contents = sqlite_session.scalars(select(HumanInputContent)).all() - assert {(content.form_id, content.message_id) for content in contents} == { - ("first-form", "message-1"), - ("second-form", "message-1"), - } - assert len(contents) == 2 + publish_spy.side_effect = complete_second_form_during_execution + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + + published = [call.args[0] for call in publish_spy.call_args_list] + completions = [ + event + for event in published + if isinstance(event, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) + ] + assert [event.form_id for event in completions] == [first.id, second.id] + assert isinstance(published[0], QueueWorkflowStartedEvent) + assert isinstance(published[-1], QueueWorkflowSucceededEvent) + assert published.index(completions[-1]) < len(published) - 1 + + +def test_waiting_form_expiring_during_resume_publishes_timeout_before_success( + sqlite_session: Session, + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + workflow_runtime: WorkflowExecutionDependencies, +) -> None: + form = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) + form.expiration_time = datetime.now(UTC) - timedelta(seconds=1) + sqlite_session.commit() + queue, publish_spy = workflow_queue + runner = _build_runner(queue, workflow_runtime) + entry = _make_paused_workflow(runner, [form], human_input_form=form) + + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + + published = [call.args[0] for call in publish_spy.call_args_list] + assert isinstance(published[-1], QueueWorkflowSucceededEvent) + timeouts = [event for event in published if isinstance(event, QueueHumanInputFormTimeoutEvent)] + assert [event.form_id for event in timeouts] == [form.id] + assert timeouts[0].node_id == "same-human-node" + assert published.index(timeouts[0]) < len(published) - 1 + sqlite_session.refresh(form) + assert form.status == HumanInputFormStatus.TIMEOUT + + +def test_chat_history_keeps_repeated_forms_by_form_id( + sqlite_session: Session, + workflow_runtime: WorkflowExecutionDependencies, + workflow_contexts: ExecutionContexts, + tool_providers: ToolProviders, +) -> None: + sqlite_session.add_all( + [ + make_app(app_id="app", tenant_id="tenant", mode=AppMode.ADVANCED_CHAT), + make_conversation( + conversation_id="conversation-1", + app_id="app", + mode=AppMode.ADVANCED_CHAT, + inputs={}, + from_source=ConversationFromSource.CONSOLE, + ), + make_message( + message_id="message-1", + app_id="app", + conversation_id="conversation-1", + inputs={}, + query="hello", + answer="", + answer_tokens=0, + message={}, + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ), + ] + ) + sqlite_session.commit() + first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + second = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + pipeline = _build_pipeline() + pipeline._workflow_tenant_id = "tenant" + pipeline._chat_records = workflow_runtime.chat_records + pipeline._message_identity = MessageIdentity("tenant", "app", "conversation-1", "message-1") + pipeline._application_generate_entity = _advanced_chat_generate_entity() + pipeline._workflow_response_converter = _build_converter( + workflow_contexts=workflow_contexts, tool_providers=tool_providers + ) + pipeline._workflow_response_converter.workflow_start_to_stream_response( + task_id="task-1", workflow_run_id="run-1", workflow_id="workflow-1", reason=WorkflowStartReason.INITIAL + ) + for form_id in (first.id, second.id, first.id): + event = QueueHumanInputFormFilledEvent( + form_id=form_id, + node_id="same-human-node", + node_type=BuiltinNodeTypes.HUMAN_INPUT, + node_title="Approval", + rendered_content="Decision: approved", + action_id="approve", + action_text="Approve", + ) + list(pipeline._handle_human_input_form_filled_event(event)) + + contents = sqlite_session.scalars(select(HumanInputContent)).all() + assert {(content.form_id, content.message_id) for content in contents} == { + (first.id, "message-1"), + (second.id, "message-1"), + } + assert len(contents) == 2 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 d1b41f1876c963..57bb5ac8cda76b 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 @@ -3,12 +3,12 @@ import pytest -from core.app.apps.base_app_generator import BaseAppGenerator -from core.app.apps.workflow.active_workflow_tasks import ( +from services.workflow.execution.adapters.workflow.active_workflow_tasks import ( active_workflow_task, get_active_workflow_task_count, reset_active_workflow_tasks, ) +from services.workflow.execution.generation_service import consume_stream @pytest.fixture(autouse=True) @@ -54,7 +54,7 @@ def response_stream() -> Generator[dict[str, str]]: worker_thread.start() assert worker_started.wait(timeout=2) - managed_stream = BaseAppGenerator._wrap_stream_with_worker_thread_join(response_stream(), worker_thread) + managed_stream = consume_stream(response_stream(), worker_thread) assert next(managed_stream) == {"event": "workflow_finished"} def finish_stream() -> None: diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py index a6f0085a80fd1a..8f674f46c38679 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py @@ -3,23 +3,18 @@ import contextlib import contextvars import json +from dataclasses import replace from types import SimpleNamespace -from unittest.mock import MagicMock, Mock +from unittest.mock import MagicMock, Mock, create_autospec import pytest from flask import Flask -from sqlalchemy import inspect from sqlalchemy.orm import Session, sessionmaker -import core.app.apps.workflow.app_generator as _workflow_app_generator -from controllers.inner_api.plugin.wraps import get_user +import services.workflow.execution.adapters.workflow.app_generator as _workflow_app_generator from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig from core.app.apps.draft_variable_saver import DraftVariableSaverFactory from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.workflow import app_generator as _app_generator_extra_app_generator_module -from core.app.apps.workflow.app_generator import WorkflowAppGenerator -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.app.entities.queue_entities import QueueWorkflowSucceededEvent from core.app.file_access import DatabaseFileAccessController @@ -32,26 +27,31 @@ from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager from core.repositories import SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository -from core.tools.tool_file_manager import ToolFileManager -from core.workflow.file_reference import build_file_reference -from core.workflow.node_runtime import resolve_dify_run_context from core.workflow.system_variables import SystemVariableKey, get_system_text from graphon.enums import WorkflowExecutionStatus from graphon.runtime import RuntimeState, VariablePool -from graphon.variable_loader import DUMMY_VARIABLE_LOADER +from graphon.variable_loader import DUMMY_VARIABLE_LOADER, VariableLoader from graphon.variables.input_entities import VariableEntity from models.enums import CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom from models.model import App, AppMode, EndUser -from models.snippet import CustomizedSnippet from models.workflow import Workflow, WorkflowKind, WorkflowNodeExecutionTriggeredFrom, WorkflowRun, WorkflowType from repositories.entities.workflow_execution import WorkflowExecutionIdentity -from services.file_request_service import FileRequestService -from services.workflow_run_agg import WorkflowRunAgg +from repositories.workflow.definition_repository import workflow_snapshot +from services.workflow.execution.adapters.node_runtime import resolve_dify_run_context +from services.workflow.execution.adapters.workflow.app_generator import WorkflowAppGenerator +from services.workflow.execution.adapters.workflow.app_queue_manager import WorkflowAppQueueManager +from services.workflow.execution.adapters.workflow.app_runner import WorkflowAppRunner +from services.workflow.execution.ports import WorkflowRuntime from tests.helpers.unit.model_factories import make_account, make_upload_file, make_workflow """SQLite-backed tests for workflow app generation and worker reload behavior.""" +def _generator(runtime: WorkflowRuntime | None = None) -> WorkflowAppGenerator: + runtime = runtime or MagicMock() + return WorkflowAppGenerator(runtime=runtime, execution_driver=runtime.execution_driver) + + def _workflow_app_generator_workflow( *, workflow_id: str = "workflow", @@ -163,7 +163,11 @@ def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: @staticmethod def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( - app: Flask, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] + app: Flask, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + *, + workflow_runtime: WorkflowRuntime, ) -> None: app_model, workflow, end_user = _persist_generator_rows(sqlite_session) workflow.graph = json.dumps( @@ -197,8 +201,10 @@ def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( queue = WorkflowAppQueueManager(entity.task_id, end_user.id, entity.invoke_from, AppMode.WORKFLOW) execution_repository, node_repository = _repositories(sqlite_session_factory, app_model, end_user) - WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run)._generate_worker( + _generator(workflow_runtime)._generate_worker( flask_app=app, + system_user_id=end_user.session_id or "", + workflow=workflow_snapshot(workflow), application_generate_entity=entity, queue_manager=queue, context=contextvars.copy_context(), @@ -215,87 +221,28 @@ def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( assert run.status == "succeeded" assert run.outputs_dict == {"caller": end_user.session_id} - @staticmethod - def test_ensure_snippet_start_node_in_worker_returns_standard_workflow_without_lookup( - sqlite_session: Session, - ) -> None: - workflow = _workflow_app_generator_workflow() - sqlite_session.add(workflow) - sqlite_session.commit() - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) - - assert result is workflow - - @staticmethod - def test_ensure_snippet_start_node_in_worker_returns_snippet_workflow_when_snippet_missing( - sqlite_session: Session, - ) -> None: - workflow = _workflow_app_generator_workflow(app_id="snippet-1", tenant_id="tenant-1", kind=WorkflowKind.SNIPPET) - other_tenant_snippet = CustomizedSnippet( - id="snippet-1", - tenant_id="tenant-2", - name="Other tenant snippet", - description="", - type="node", - ) - sqlite_session.add_all([workflow, other_tenant_snippet]) - sqlite_session.commit() - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) - - assert result is workflow - - @staticmethod - def test_ensure_snippet_start_node_in_worker_applies_snippet_start_injection( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ) -> None: - workflow = _workflow_app_generator_workflow(app_id="snippet-1", tenant_id="tenant-1", kind=WorkflowKind.SNIPPET) - snippet = CustomizedSnippet( - id="snippet-1", - tenant_id="tenant-1", - name="Matching snippet", - description="", - type="node", - ) - sqlite_session.add_all([workflow, snippet]) - sqlite_session.commit() - ensure_start_node = MagicMock(return_value=workflow) - monkeypatch.setattr( - "services.snippet_generate_service.SnippetGenerateService.ensure_start_node_for_worker", - ensure_start_node, - ) - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) - - assert result is workflow - ensure_start_node.assert_called_once_with(workflow, snippet) - @staticmethod def test_generate_includes_parent_trace_context_in_extras( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, *, workflow_runtime: WorkflowRuntime ) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = _generator(workflow_runtime) app, workflow, end_user = _persist_generator_rows(sqlite_session) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerator._bind_file_access_scope", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppGenerator._bind_file_access_scope", lambda *args, **kwargs: contextlib.nullcontext(), ) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppConfigManager.get_app_config", lambda *args, **kwargs: _app_config(app, workflow), ) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.file_factory.build_from_mappings", lambda *args, **kwargs: [] + "services.workflow.execution.adapters.workflow.app_generator.file_factory.build_from_mappings", + lambda *args, **kwargs: [], ) repository_tenant_ids: dict[str, str] = {} - workflow_execution_factory = ( - _workflow_app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository - ) - workflow_node_execution_factory = ( - _workflow_app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository - ) + workflow_execution_factory = workflow_runtime.execution_writer + workflow_node_execution_factory = workflow_runtime.node_writer def create_workflow_execution_repository(**kwargs): repository_tenant_ids["workflow"] = kwargs["tenant_id"] @@ -305,16 +252,10 @@ def create_workflow_node_execution_repository(**kwargs): repository_tenant_ids["node"] = kwargs["tenant_id"] return workflow_node_execution_factory(**kwargs) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - create_workflow_execution_repository, - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - create_workflow_node_execution_repository, - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.db", SimpleNamespace(engine=sqlite_session.get_bind()) + generator._runtime = replace( + workflow_runtime, + execution_writer=create_workflow_execution_repository, + node_writer=create_workflow_node_execution_repository, ) monkeypatch.setattr(generator, "_prepare_user_inputs", lambda *, user_inputs, **kwargs: user_inputs) @@ -349,6 +290,8 @@ def fake_generate(**kwargs): assert result == {"data": {}} application_generate_entity = captured["application_generate_entity"] assert isinstance(application_generate_entity, WorkflowAppGenerateEntity) + assert application_generate_entity.user_id == end_user.id + assert application_generate_entity.user_from == UserFrom.ACCOUNT extras = application_generate_entity.extras assert extras["external_trace_id"] == "trace-1" assert extras["parent_trace_context"].model_dump() == { @@ -365,8 +308,10 @@ def test_resume_delegates_to_generate( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], + *, + workflow_runtime: WorkflowRuntime, ) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = _generator(workflow_runtime) app, workflow, end_user = _persist_generator_rows(sqlite_session) mock_generate = MagicMock(return_value="ok") monkeypatch.setattr(generator, "_generate", mock_generate) @@ -389,7 +334,7 @@ def test_resume_delegates_to_generate( result = generator.resume( app_model=app, - workflow=workflow, + workflow=workflow_snapshot(workflow), user=end_user, application_generate_entity=application_generate_entity, graph_runtime_state=runtime_state, @@ -412,24 +357,26 @@ def test_generate_forwards_pause_configuration_and_state( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], + *, + workflow_runtime: WorkflowRuntime, ) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = _generator(workflow_runtime) app, workflow, end_user = _persist_generator_rows(sqlite_session) queue_manager = MagicMock() monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppQueueManager", MagicMock(return_value=queue_manager), ) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", MagicMock(return_value="converted"), ) monkeypatch.setattr(WorkflowAppGenerator, "_handle_response", MagicMock(return_value="response")) draft_factory_tenant_ids: list[str] = [] - get_draft_var_saver_factory = WorkflowAppGenerator._get_draft_var_saver_factory + get_draft_var_saver_factory = generator._get_draft_var_saver_factory def get_recording_draft_var_saver_factory( invoke_from: InvokeFrom, account: EndUser, *, tenant_id: str @@ -443,9 +390,6 @@ def get_recording_draft_var_saver_factory( staticmethod(get_recording_draft_var_saver_factory), ) - db_session = sqlite_session_factory() - monkeypatch.setattr(_workflow_app_generator.db, "session", db_session) - worker_kwargs: dict[str, object] = {} class DummyThread: @@ -463,7 +407,7 @@ def join(self, timeout): def is_alive(self): return False - monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", DummyThread) + monkeypatch.setattr("services.workflow.execution.adapters.workflow.app_generator.threading.Thread", DummyThread) application_generate_entity = _generate_entity(app, workflow, end_user) graph_runtime_state = _runtime_state() @@ -475,7 +419,7 @@ def is_alive(self): with flask_app.app_context(): result = generator._generate( app_model=app, - workflow=workflow, + workflow=workflow_snapshot(workflow), user=end_user, application_generate_entity=application_generate_entity, invoke_from=InvokeFrom.SERVICE_API, @@ -508,8 +452,10 @@ def test_execution_grants_default_files_to_worker_and_response( sqlite_session_factory: sessionmaker[Session], resuming: bool, invoke_from: InvokeFrom, + *, + workflow_runtime: WorkflowRuntime, ) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = _generator(workflow_runtime) app, workflow, end_user = _persist_generator_rows(sqlite_session) default_file, unrelated_file = [ make_upload_file( @@ -566,7 +512,7 @@ def assert_default_file_access() -> None: queue_manager = MagicMock() monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppQueueManager", MagicMock(return_value=queue_manager), ) @@ -576,13 +522,10 @@ def handle_response(**_kwargs): monkeypatch.setattr(generator, "_handle_response", handle_response) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", MagicMock(side_effect=lambda response, invoke_from: response), ) - db_session = sqlite_session_factory() - monkeypatch.setattr(_workflow_app_generator.db, "session", db_session) - runner_instance = MagicMock() runner_instance.prepare.return_value = None @@ -595,7 +538,7 @@ def runner_ctor(**kwargs): return runner_instance monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppRunner", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppRunner", MagicMock(side_effect=runner_ctor), ) @@ -615,7 +558,9 @@ def join(self, timeout): def is_alive(self): return False - monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", ImmediateThread) + monkeypatch.setattr( + "services.workflow.execution.adapters.workflow.app_generator.threading.Thread", ImmediateThread + ) pause_config = PauseStateLayerConfig( session_factory=sqlite_session_factory, @@ -630,7 +575,7 @@ def is_alive(self): kwargs = { "app_model": app, - "workflow": workflow, + "workflow": workflow_snapshot(workflow), "user": end_user, "application_generate_entity": application_generate_entity, "graph_runtime_state": runtime_state, @@ -647,7 +592,6 @@ def is_alive(self): TENANT_ID = "00000000-0000-0000-0000-000000000001" -OTHER_TENANT_ID = "00000000-0000-0000-0000-000000000002" APP_ID = "00000000-0000-0000-0000-000000000003" WORKFLOW_ID = "00000000-0000-0000-0000-000000000004" END_USER_ID = "00000000-0000-0000-0000-000000000005" @@ -741,24 +685,6 @@ def _persist_end_user(session: Session) -> EndUser: return end_user -def _persist_snippet( - session: Session, - *, - snippet_id: str, - tenant_id: str = TENANT_ID, -) -> CustomizedSnippet: - snippet = CustomizedSnippet( - id=snippet_id, - tenant_id=tenant_id, - name="Snippet", - description=None, - type="node", - ) - session.add(snippet) - session.commit() - return snippet - - class TestAppGeneratorExtraCases: @staticmethod @pytest.fixture(autouse=True) @@ -768,32 +694,19 @@ def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(TraceQueueManager, "start_timer", lambda self: None) monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - @staticmethod - @pytest.fixture - def sqlite_generator_session( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine, - sqlite_session: Session, - ) -> Session: - """Expose the shared SQLite engine and real Session through the Flask extension boundary.""" - monkeypatch.setattr( - _app_generator_extra_app_generator_module, - "db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), - ) - return sqlite_session - @pytest.mark.parametrize("resuming", ["fresh", "new-snapshot", "legacy-snapshot"]) @pytest.mark.parametrize("is_account", [True, False], ids=["account", "end-user"]) - def test_service_api_preserves_caller_identity_and_file_access( + def test_service_api_preserves_caller_identity_across_resume( self, monkeypatch: pytest.MonkeyPatch, - sqlite_generator_session: Session, + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], resuming: str, is_account: bool, + *, + workflow_runtime: WorkflowRuntime, ) -> None: - app_model, workflow, end_user = _persist_generator_rows(sqlite_generator_session) + app_model, workflow, end_user = _persist_generator_rows(sqlite_session) account = make_account(account_id="scheduled-account") user = account if is_account else end_user expected_user_from = UserFrom.ACCOUNT if is_account and resuming != "legacy-snapshot" else UserFrom.END_USER @@ -819,7 +732,7 @@ def test_service_api_preserves_caller_identity_and_file_access( observed: list[tuple[str, UserFrom, str]] = [] def inspect_execution( - runner: WorkflowBasedAppRunner, + runner: WorkflowAppRunner, _pause_config: PauseStateLayerConfig | None, _executions: WorkflowExecutionRepository, _nodes: WorkflowNodeExecutionRepository, @@ -831,74 +744,23 @@ def inspect_execution( prepared.entry.graph_engine.runtime_state.variable_pool, SystemVariableKey.USER_ID ) assert context.user_from == expected_user_from - # Old scheduled snapshots retain the Agent upload's EndUser remapping, - # including its known mismatch with download's original Account id. - legacy_scheduled = resuming == "legacy-snapshot" and is_account - owner_id = get_user(context.tenant_id, context.user_id).id if legacy_scheduled else context.user_id - tool_file = ToolFileManager().create_file_by_raw( - user_id=owner_id, - tenant_id=context.tenant_id, - conversation_id=None, - file_binary=b"scheduled report", - mimetype="text/plain", - filename="report.txt", - ) - mapping = {"transfer_method": "tool_file", "reference": build_file_reference(record_id=tool_file.id)} - service = FileRequestService() - if legacy_scheduled: - assert owner_id != context.user_id - with pytest.raises(ValueError, match="not found"): - service.request_download( - tenant_id=context.tenant_id, - user_id=context.user_id, - user_from=context.user_from, - invoke_from=context.invoke_from, - file_mapping=mapping, - ) - observed.append((context.user_id, context.user_from, system_user_id or "")) - return - result = service.request_download( - tenant_id=context.tenant_id, - user_id=context.user_id, - user_from=context.user_from, - invoke_from=context.invoke_from, - file_mapping=mapping, - ) - assert result.filename == "report.txt" - assert result.size == len(b"scheduled report") - assert tool_file.id in result.download_uri - for tenant_id, user_id, user_from in [ - (context.tenant_id, "another-end-user", UserFrom.END_USER), - ("another-tenant", context.user_id, context.user_from), - ]: - with pytest.raises(ValueError, match="not found"): - service.request_download( - tenant_id=tenant_id, - user_id=user_id, - user_from=user_from, - invoke_from=InvokeFrom.SERVICE_API, - file_mapping=mapping, - ) - observed.append((context.user_id, context.user_from, system_user_id or "")) - generator = WorkflowAppGenerator(execution_driver=inspect_execution) + generator = WorkflowAppGenerator(runtime=workflow_runtime, execution_driver=inspect_execution) monkeypatch.setattr(generator, "_handle_response", lambda **kwargs: {}) monkeypatch.setattr( _workflow_app_generator.WorkflowAppGenerateResponseConverter, "convert", lambda response, invoke_from: response, ) - monkeypatch.setattr(generator, "_restore_workflow_run_graph", lambda **kwargs: None) - monkeypatch.setattr("core.tools.tool_file_manager.storage.save", lambda *args: None) - monkeypatch.setattr("controllers.inner_api.plugin.wraps.db", _workflow_app_generator.db) + monkeypatch.setattr(workflow_runtime.contexts, "restore_graph", lambda workflow, _run_id: workflow) if resuming == "fresh": generator.generate( app_model=app_model, workflow=workflow, user=user, execution_identity=WorkflowExecutionIdentity(user_id=user.id, user_from=expected_user_from), - args={"inputs": {}}, + args={"inputs": {}, "files": []}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, ) @@ -906,7 +768,7 @@ def inspect_execution( assert runtime_state is not None generator.resume( app_model=app_model, - workflow=workflow, + workflow=workflow_snapshot(workflow), user=user, application_generate_entity=entity, workflow_execution_repository=execution_repository, @@ -916,11 +778,14 @@ def inspect_execution( assert observed == [(user.id, expected_user_from, expected_system_user_id)] class TestWorkflowAppGeneratorValidation: - @pytest.mark.usefixtures("sqlite_generator_session") def test_generate_stream_joins_worker_after_response_exhaustion( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + *, + workflow_runtime: WorkflowRuntime, ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = _generator(workflow_runtime) app = _persist_app(sqlite_session) workflow = _persist_workflow(sqlite_session) user = _persist_end_user(sqlite_session) @@ -951,19 +816,24 @@ def response_stream(): monkeypatch.setattr(generator, "_bind_file_access_scope", lambda **kwargs: contextlib.nullcontext()) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppQueueManager", lambda **kwargs: SimpleNamespace(**kwargs), ) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.current_app", + "services.workflow.execution.adapters.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( + "services.workflow.execution.adapters.workflow.app_generator.contextvars.copy_context", lambda: "ctx" + ) + monkeypatch.setattr( + "services.workflow.execution.adapters.workflow.app_generator.threading.Thread", + lambda **kwargs: worker_thread, + ) 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", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", lambda response, invoke_from: response, ) @@ -983,57 +853,8 @@ def response_stream(): 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, - unbound_session: Session, - ): - workflow = _app_generator_extra_workflow() - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( - session=unbound_session, - workflow=workflow, - ) - - assert result is workflow - - def test_ensure_snippet_start_node_returns_original_when_snippet_is_from_another_tenant( - self, - sqlite_session: Session, - ): - workflow = _persist_workflow(sqlite_session, kind=WorkflowKind.SNIPPET) - _persist_snippet(sqlite_session, snippet_id=APP_ID, tenant_id=OTHER_TENANT_ID) - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( - session=sqlite_session, - workflow=workflow, - ) - - assert result is workflow - - def test_ensure_snippet_start_node_delegates_when_snippet_exists( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ): - workflow = _persist_workflow(sqlite_session, kind=WorkflowKind.SNIPPET) - snippet = _persist_snippet(sqlite_session, snippet_id=APP_ID) - injected_workflow = _app_generator_extra_workflow(workflow_id="workflow-injected") - ensure_start_node = Mock(return_value=injected_workflow) - monkeypatch.setattr( - "services.snippet_generate_service.SnippetGenerateService.ensure_start_node_for_worker", - ensure_start_node, - ) - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( - session=sqlite_session, - workflow=workflow, - ) - - assert result is injected_workflow - ensure_start_node.assert_called_once_with(workflow, snippet) - - def test_single_iteration_generate_validates_args(self, sqlite_session: Session): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_single_iteration_generate_validates_args(self, *, workflow_runtime: WorkflowRuntime): + generator = _generator(workflow_runtime) with pytest.raises(ValueError, match="node_id is required"): generator.single_iteration_generate( @@ -1044,7 +865,6 @@ def test_single_iteration_generate_validates_args(self, sqlite_session: Session) execution_identity=WorkflowExecutionIdentity(user_id=_end_user().id, user_from=UserFrom.END_USER), args={"inputs": {}}, streaming=False, - session=sqlite_session, ) with pytest.raises(ValueError, match="inputs is required"): @@ -1056,11 +876,10 @@ def test_single_iteration_generate_validates_args(self, sqlite_session: Session) execution_identity=WorkflowExecutionIdentity(user_id=_end_user().id, user_from=UserFrom.END_USER), args={}, streaming=False, - session=sqlite_session, ) - def test_single_loop_generate_validates_args(self, sqlite_session: Session): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_single_loop_generate_validates_args(self, *, workflow_runtime: WorkflowRuntime): + generator = _generator(workflow_runtime) with pytest.raises(ValueError, match="node_id is required"): generator.single_loop_generate( @@ -1069,18 +888,24 @@ def test_single_loop_generate_validates_args(self, sqlite_session: Session): node_id="", user=_end_user(), execution_identity=WorkflowExecutionIdentity(user_id=_end_user().id, user_from=UserFrom.END_USER), - args=SimpleNamespace(inputs={}), + args={"inputs": {}}, streaming=False, - session=sqlite_session, ) - @pytest.mark.usefixtures("sqlite_generator_session") def test_single_iteration_generate_includes_trace_session_id_in_extras( self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, + *, + workflow_runtime: WorkflowRuntime, ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + loader = create_autospec(VariableLoader, instance=True) + loader_factory = Mock(return_value=loader) + generator = WorkflowAppGenerator( + draft_variable_loader=loader_factory, + runtime=workflow_runtime, + execution_driver=workflow_runtime.execution_driver, + ) app = _persist_app(sqlite_session) workflow = _persist_workflow(sqlite_session) user = _persist_end_user(sqlite_session) @@ -1093,31 +918,10 @@ def test_single_iteration_generate_includes_trace_session_id_in_extras( workflow_id=WORKFLOW_ID, ) captured: dict[str, object] = {} - repository_session_makers: list[sessionmaker[Session]] = [] - draft_sessions: list[Session] = [] - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppConfigManager.get_app_config", lambda **kwargs: app_config, ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DraftVarLoader", lambda **kwargs: SimpleNamespace() - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowDraftVariableService", - lambda session: ( - draft_sessions.append(session) - or SimpleNamespace(prefill_conversation_variable_default_values=lambda *args, **kwargs: None) - ), - ) monkeypatch.setattr(generator, "_generate", lambda **kwargs: captured.update(kwargs) or {"ok": True}) generator.single_iteration_generate( @@ -1128,22 +932,27 @@ def test_single_iteration_generate_includes_trace_session_id_in_extras( execution_identity=WorkflowExecutionIdentity(user_id=user.id, user_from=UserFrom.END_USER), args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, streaming=False, - session=sqlite_session, ) assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" - assert len(repository_session_makers) == 2 - assert all(factory.kw["bind"] is sqlite_session.get_bind() for factory in repository_session_makers) - assert len(draft_sessions) == 1 - assert draft_sessions[0] is sqlite_session + assert generator._runtime is workflow_runtime + loader_factory.assert_called_once_with(workflow, user.id) + assert captured["variable_loader"] is loader - @pytest.mark.usefixtures("sqlite_generator_session") def test_single_loop_generate_includes_trace_session_id_in_extras( self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, + *, + workflow_runtime: WorkflowRuntime, ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + loader = create_autospec(VariableLoader, instance=True) + loader_factory = Mock(return_value=loader) + generator = WorkflowAppGenerator( + draft_variable_loader=loader_factory, + runtime=workflow_runtime, + execution_driver=workflow_runtime.execution_driver, + ) app = _persist_app(sqlite_session) workflow = _persist_workflow(sqlite_session) user = _persist_end_user(sqlite_session) @@ -1156,31 +965,10 @@ def test_single_loop_generate_includes_trace_session_id_in_extras( workflow_id=WORKFLOW_ID, ) captured: dict[str, object] = {} - repository_session_makers: list[sessionmaker[Session]] = [] - draft_sessions: list[Session] = [] - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppConfigManager.get_app_config", lambda **kwargs: app_config, ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DraftVarLoader", lambda **kwargs: SimpleNamespace() - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowDraftVariableService", - lambda session: ( - draft_sessions.append(session) - or SimpleNamespace(prefill_conversation_variable_default_values=lambda *args, **kwargs: None) - ), - ) monkeypatch.setattr(generator, "_generate", lambda **kwargs: captured.update(kwargs) or {"ok": True}) generator.single_loop_generate( @@ -1189,16 +977,12 @@ def test_single_loop_generate_includes_trace_session_id_in_extras( node_id="node-2", user=user, execution_identity=WorkflowExecutionIdentity(user_id=user.id, user_from=UserFrom.END_USER), - args=SimpleNamespace(inputs={"foo": "bar"}, trace_session_id="session-1"), + args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, streaming=False, - session=sqlite_session, ) assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" - assert len(repository_session_makers) == 2 - assert all(factory.kw["bind"] is sqlite_session.get_bind() for factory in repository_session_makers) - assert len(draft_sessions) == 1 - assert draft_sessions[0] is sqlite_session + assert generator._runtime is workflow_runtime with pytest.raises(ValueError, match="inputs is required"): generator.single_loop_generate( @@ -1207,14 +991,18 @@ def test_single_loop_generate_includes_trace_session_id_in_extras( node_id="node", user=_end_user(), execution_identity=WorkflowExecutionIdentity(user_id=_end_user().id, user_from=UserFrom.END_USER), - args=SimpleNamespace(inputs=None), + args={"inputs": None}, streaming=False, - session=sqlite_session, ) + loader_factory.assert_called_once_with(workflow, user.id) + assert captured["variable_loader"] is loader + class TestWorkflowAppGeneratorHandleResponse: - def test_handle_response_closed_file_raises_stopped(self, monkeypatch: pytest.MonkeyPatch): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_handle_response_closed_file_raises_stopped( + self, monkeypatch: pytest.MonkeyPatch, *, workflow_runtime: WorkflowRuntime + ): + generator = _generator(workflow_runtime) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", @@ -1247,7 +1035,7 @@ def process(self): raise ValueError("I/O operation on closed file.") monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerateTaskPipeline", + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppGenerateTaskPipeline", _Pipeline, ) @@ -1263,12 +1051,18 @@ def process(self): class TestWorkflowAppGeneratorGenerate: @pytest.fixture - def generation(self, monkeypatch: pytest.MonkeyPatch, sqlite_generator_session: Session): + def generation( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + *, + workflow_runtime: WorkflowRuntime, + ): """Keep input preparation real and stop at the execution boundary.""" - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app = _persist_app(sqlite_generator_session) - workflow = _persist_workflow(sqlite_generator_session) - user = _persist_end_user(sqlite_generator_session) + generator = _generator(workflow_runtime) + app = _persist_app(sqlite_session) + workflow = _persist_workflow(sqlite_session) + user = _persist_end_user(sqlite_session) execute = Mock(return_value={"ok": True}) monkeypatch.setattr(generator, "_generate", execute) return generator, app, workflow, user, execute @@ -1360,8 +1154,10 @@ def test_generate_prepares_start_inputs(self, generation, root_node_id): execute.assert_not_called() class TestWorkflowAppGeneratorResume: - def test_resume_restores_trace_manager_when_missing(self, monkeypatch: pytest.MonkeyPatch): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_resume_restores_trace_manager_when_missing( + self, monkeypatch: pytest.MonkeyPatch, *, workflow_runtime: WorkflowRuntime + ): + generator = _generator(workflow_runtime) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", app_id="app", @@ -1394,7 +1190,7 @@ def test_resume_restores_trace_manager_when_missing(self, monkeypatch: pytest.Mo }, ) monkeypatch.setattr( - "core.app.apps.workflow.app_generator.TraceQueueManager", + "services.workflow.execution.adapters.workflow.app_generator.TraceQueueManager", dummy_trace_queue_manager, ) captured_entity: WorkflowAppGenerateEntity | None = None @@ -1423,8 +1219,10 @@ def _fake_generate(**kwargs): assert trace_manager.app_id == APP_ID assert trace_manager.user_id == "session-id" - def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.MonkeyPatch): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + def test_resume_preserves_existing_trace_manager( + self, monkeypatch: pytest.MonkeyPatch, *, workflow_runtime: WorkflowRuntime + ): + generator = _generator(workflow_runtime) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", app_id="app", @@ -1476,11 +1274,14 @@ def test_generate_worker_uses_end_user_session_for_external_invocation( self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, + *, + workflow_runtime: WorkflowRuntime, ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + generator = _generator(workflow_runtime) _persist_app(sqlite_session) - _persist_workflow(sqlite_session) + workflow = _persist_workflow(sqlite_session) _persist_end_user(sqlite_session) + sqlite_session.expunge(workflow) runner_kwargs = {} @@ -1492,12 +1293,14 @@ def prepare(self): return None monkeypatch.setattr( - "core.app.apps.workflow.app_generator.preserve_flask_contexts", + "services.workflow.execution.adapters.workflow.app_generator.preserve_flask_contexts", lambda flask_app, context_vars: contextlib.nullcontext(), ) - monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppRunner", _Runner) - restore_workflow_run_graph = Mock() - monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph) + monkeypatch.setattr( + "services.workflow.execution.adapters.workflow.app_generator.WorkflowAppRunner", _Runner + ) + restore_workflow_run_graph = Mock(side_effect=lambda workflow, _run_id: workflow) + monkeypatch.setattr(generator._runtime.contexts, "restore_graph", restore_workflow_run_graph) app_config = WorkflowUIBasedAppConfig( tenant_id=TENANT_ID, @@ -1523,6 +1326,8 @@ def prepare(self): ) generator._generate_worker( + system_user_id="session-id", + workflow=workflow_snapshot(workflow), flask_app=SimpleNamespace(), application_generate_entity=application_generate_entity, queue_manager=SimpleNamespace(), @@ -1534,9 +1339,7 @@ def prepare(self): ) assert runner_kwargs["system_user_id"] == "session-id" + assert runner_kwargs["runtime"] is workflow_runtime + assert "execution_driver" not in runner_kwargs restore_workflow_run_graph.assert_called_once() - restore_kwargs = restore_workflow_run_graph.call_args.kwargs - assert isinstance(restore_kwargs["session"], Session) - assert restore_kwargs["workflow"] is runner_kwargs["workflow"] - assert restore_kwargs["workflow_run_id"] == "run-id" - assert inspect(runner_kwargs["workflow"]).detached is True + restore_workflow_run_graph.assert_called_once_with(runner_kwargs["workflow"], "run-id") diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py b/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py index 9b5e700728aec4..acc1f3e50359dc 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py @@ -6,12 +6,12 @@ import pytest from core.app.apps.execution_coordinator import send_abort_command -from core.app.apps.workflow.command_channels import ( +from graphon.engine.command import AbortCommand, PauseCommand, RedisChannel +from services.workflow.execution.adapters.workflow.command_channels import ( CelerySignalCommandChannel, CombinedCommandChannel, StopFlagCommandChannel, ) -from graphon.engine.command import AbortCommand, PauseCommand, RedisChannel class _CommandChannelStub: @@ -124,10 +124,14 @@ def test_celery_signal_command_channel_emits_abort_once_per_instance() -> None: def test_stop_flag_command_channel_emits_abort_when_flag_is_set() -> None: channel = StopFlagCommandChannel(task_id="task-1", abort_reason="User requested stop") - with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=False): + with patch( + "services.workflow.execution.adapters.workflow.command_channels.is_app_task_stop_flag_set", return_value=False + ): assert channel.fetch_commands() == [] - with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=True): + with patch( + "services.workflow.execution.adapters.workflow.command_channels.is_app_task_stop_flag_set", return_value=True + ): commands = channel.fetch_commands() assert len(commands) == 1 @@ -138,6 +142,8 @@ def test_stop_flag_command_channel_emits_abort_when_flag_is_set() -> None: def test_stop_flag_command_channel_emits_abort_once_per_instance() -> None: channel = StopFlagCommandChannel(task_id="task-1") - with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=True): + with patch( + "services.workflow.execution.adapters.workflow.command_channels.is_app_task_stop_flag_set", return_value=True + ): assert len(channel.fetch_commands()) == 1 assert channel.fetch_commands() == [] diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py index d0568e7cabd1c8..6bbbd6b9f2fdf4 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py @@ -1,17 +1,15 @@ from __future__ import annotations -import json -import time from types import SimpleNamespace from unittest.mock import MagicMock import pytest -from sqlalchemy import select +from sqlalchemy import create_engine, event, select from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.pool import QueuePool from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline +from core.app.apps.draft_variable_saver import NoopDraftVariableSaver from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.app.entities.queue_entities import ( QueueAgentLogEvent, @@ -52,17 +50,21 @@ ) from core.base.tts.app_generator_tts_publisher import AudioTrunk from core.workflow.system_variables import build_system_variables, system_variables_to_mapping +from extensions.application_services.workflow import build_workflow_execution_dependencies from extensions.ext_redis import RedisClientWrapper from graphon.entities import WorkflowStartReason from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.runtime import RuntimeState, VariablePool from libs.datetime_utils import naive_utc_now -from models.account import Account +from models.base import TypeBase from models.enums import CreatorUserRole from models.model import AppMode, EndUser -from models.workflow import Workflow, WorkflowAppLog, WorkflowType -from tests.helpers.unit.model_factories import make_end_user, make_workflow +from models.workflow import Workflow, WorkflowAppLog +from repositories.workflow.definition_repository import workflow_snapshot +from services.workflow.execution.adapters.workflow.app_queue_manager import WorkflowAppQueueManager +from services.workflow.execution.adapters.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline +from tests.helpers.unit.model_factories import make_account, make_end_user, make_workflow from tests.workflow_test_utils import build_test_variable_pool @@ -75,51 +77,150 @@ def _build_workflow_app_config() -> WorkflowUIBasedAppConfig: ) -def _build_generate_entity(run_id: str) -> WorkflowAppGenerateEntity: - return WorkflowAppGenerateEntity( - user_from=UserFrom.END_USER, +@pytest.fixture(autouse=True) +def bind_queue_redis(monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None: + monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0]) + + +@pytest.fixture +def log_store(): + # This database is deliberately separate from the application's global database. + engine = create_engine("sqlite://", poolclass=QueuePool) + TypeBase.metadata.create_all(engine, tables=[WorkflowAppLog.__table__]) + sessions = sessionmaker(engine, expire_on_commit=False) + try: + yield sessions, build_workflow_execution_dependencies(sessions) + finally: + engine.dispose() + + +def build_pipeline(runtime, invoke_from, user, *, tool_providers, run_id: str = "run-id"): + workflow = make_workflow(workflow_id="workflow-id", app_id="app-id", tenant_id="tenant-id", features="{}") + entity = WorkflowAppGenerateEntity( + user_from=UserFrom.END_USER if isinstance(user, EndUser) else UserFrom.ACCOUNT, task_id="task-id", - app_config=_build_workflow_app_config(), + app_config=WorkflowUIBasedAppConfig( + tenant_id=workflow.tenant_id, + app_id=workflow.app_id, + app_mode=AppMode.WORKFLOW, + workflow_id=workflow.id, + ), inputs={}, files=[], - user_id="user-id", + user_id=user.id, stream=False, - invoke_from=InvokeFrom.SERVICE_API, + invoke_from=invoke_from, workflow_execution_id=run_id, ) - - -def _build_runtime_state(run_id: str) -> RuntimeState: - variable_pool = build_test_variable_pool(variables=build_system_variables(workflow_execution_id=run_id)) - return RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=time.perf_counter()) - - -def _build_pipeline(run_id: str, unbound_session_factory: sessionmaker[Session]) -> WorkflowAppGenerateTaskPipeline: - queue_manager = WorkflowAppQueueManager( - task_id="task-id", user_id="user-id", invoke_from=InvokeFrom.SERVICE_API, app_mode=AppMode.WORKFLOW + state = RuntimeState( + workflow_id=workflow.id, + variable_pool=build_test_variable_pool(variables=build_system_variables(workflow_execution_id=run_id)), + start_at=0, ) - queue_manager.graph_runtime_state = _build_runtime_state(run_id) - workflow = Workflow( - id="workflow-id", - tenant_id="tenant-id", - app_id="app-id", - type=WorkflowType.WORKFLOW, - version=Workflow.VERSION_DRAFT, - graph="{}", - _features=json.dumps({}), - created_by="user-id", + queue_manager = WorkflowAppQueueManager( + task_id=entity.task_id, + user_id=user.id, + invoke_from=invoke_from, + app_mode=AppMode.WORKFLOW, ) - user = Account(name="user", email="user@example.com") - pipeline = WorkflowAppGenerateTaskPipeline( - application_generate_entity=_build_generate_entity(run_id), - workflow=workflow, + queue_manager.graph_runtime_state = state + return WorkflowAppGenerateTaskPipeline( + contexts=runtime.contexts, + logs=runtime.logs, + application_generate_entity=entity, + workflow=workflow_snapshot(workflow), queue_manager=queue_manager, user=user, stream=False, - draft_var_saver_factory=MagicMock(), + draft_var_saver_factory=lambda **_kwargs: NoopDraftVariableSaver(), + tool_providers=tool_providers, + ) + + +@pytest.mark.parametrize( + ("invoke_from", "created_from"), + [ + (InvokeFrom.SERVICE_API, "service-api"), + (InvokeFrom.OPENAPI, "openapi"), + (InvokeFrom.EXPLORE, "installed-app"), + (InvokeFrom.WEB_APP, "web-app"), + ], +) +@pytest.mark.parametrize("account", [True, False]) +def test_start_commits_log_to_injected_database_before_response( + log_store, invoke_from, created_from, account, *, tool_providers +): + sessions, runtime = log_store + user = make_account(account_id="user-id") if account else make_end_user(end_user_id="user-id") + pipeline = build_pipeline(runtime, invoke_from, user, tool_providers=tool_providers) + responses = pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent(reason=WorkflowStartReason.INITIAL)) + + assert next(responses).workflow_run_id == "run-id" + assert sessions.kw["bind"].pool.checkedout() == 0 + with sessions() as session: + log = session.scalars(select(WorkflowAppLog)).one() + assert (log.tenant_id, log.app_id, log.workflow_id, log.workflow_run_id) == ( + "tenant-id", + "app-id", + "workflow-id", + "run-id", + ) + assert log.created_from == created_from + assert log.created_by == user.id + assert log.created_by_role == (CreatorUserRole.ACCOUNT if account else CreatorUserRole.END_USER) + assert list(responses) == [] + + list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent(reason=WorkflowStartReason.RESUMPTION))) + with sessions() as session: + assert session.scalars(select(WorkflowAppLog)).one().id == log.id + + +@pytest.mark.parametrize( + "invoke_from", [InvokeFrom.DEBUGGER, InvokeFrom.TRIGGER, InvokeFrom.PUBLISHED_PIPELINE, InvokeFrom.VALIDATION] +) +def test_debug_and_pipeline_starts_do_not_create_application_logs(log_store, invoke_from, *, tool_providers): + sessions, runtime = log_store + pipeline = build_pipeline(runtime, invoke_from, make_account(account_id="user-id"), tool_providers=tool_providers) + assert len(list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent()))) == 1 + with sessions() as session: + assert session.scalar(select(WorkflowAppLog)) is None + + +def test_log_failure_rolls_back_and_releases_connection_for_retry(log_store, *, tool_providers): + sessions, runtime = log_store + pipeline = build_pipeline( + runtime, InvokeFrom.SERVICE_API, make_account(account_id="user-id"), tool_providers=tool_providers + ) + + def fail_after_insert(session, _flush_context): + assert session.scalar(select(WorkflowAppLog)) is not None + raise RuntimeError("commit failed") + + event.listen(sessions, "after_flush", fail_after_insert) + try: + with pytest.raises(RuntimeError, match="commit failed"): + list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) + finally: + event.remove(sessions, "after_flush", fail_after_insert) + + assert sessions.kw["bind"].pool.checkedout() == 0 + with sessions() as session: + assert session.scalar(select(WorkflowAppLog)) is None + + assert len(list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent()))) == 1 + with sessions() as session: + assert session.scalars(select(WorkflowAppLog)).one().workflow_run_id == "run-id" + + +def _build_pipeline(run_id: str, unbound_session_factory: sessionmaker[Session]): + runtime = build_workflow_execution_dependencies(unbound_session_factory) + return build_pipeline( + runtime, + InvokeFrom.WEB_APP, + make_account(account_id="user-id"), + tool_providers=runtime.tool_providers, + run_id=run_id, ) - pipeline._database_session = unbound_session_factory - return pipeline class TestGenerateTaskPipelineCases: @@ -189,7 +290,7 @@ def _make_pipeline(): extras={}, call_depth=0, ) - workflow = _workflow() + workflow = workflow_snapshot(_workflow()) user = _end_user() pipeline = WorkflowAppGenerateTaskPipeline( @@ -199,6 +300,9 @@ def _make_pipeline(): user=user, stream=False, draft_var_saver_factory=lambda **kwargs: None, + contexts=MagicMock(), + logs=MagicMock(), + tool_providers=MagicMock(), ) return pipeline @@ -306,7 +410,7 @@ def test_handle_error_event(self): assert isinstance(responses[0], ValueError) - def test_handle_workflow_started_event_sets_run_id(self, monkeypatch: pytest.MonkeyPatch, sqlite_engine): + def test_handle_workflow_started_event_sets_run_id(self, monkeypatch: pytest.MonkeyPatch): pipeline = _make_pipeline() pipeline._graph_runtime_state = RuntimeState( workflow_id="test-workflow", @@ -317,10 +421,6 @@ def test_handle_workflow_started_event_sets_run_id(self, monkeypatch: pytest.Mon ) pipeline._workflow_response_converter.workflow_start_to_stream_response = lambda **kwargs: "started" - monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.db", - SimpleNamespace(engine=sqlite_engine), - ) monkeypatch.setattr(pipeline, "_save_workflow_app_log", lambda **kwargs: None) responses = list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) @@ -460,18 +560,16 @@ def test_handle_stop_event_yields_finish(self): assert responses == ["finish"] - @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) - def test_save_workflow_app_log_created_from(self, sqlite_session: Session): + def test_save_workflow_app_log_created_from(self): pipeline = _make_pipeline() pipeline._application_generate_entity.invoke_from = InvokeFrom.SERVICE_API pipeline._user_id = "user" - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id") - sqlite_session.flush() + pipeline._save_workflow_app_log(workflow_run_id="run-id") - saved_log = sqlite_session.scalar(select(WorkflowAppLog)) - assert saved_log is not None - assert saved_log.workflow_run_id == "run-id" - assert saved_log.created_from == "service-api" + pipeline._logs.record.assert_called_once() + saved_log = pipeline._logs.record.call_args.kwargs + assert saved_log["workflow_run_id"] == "run-id" + assert saved_log["created_from"] == "service-api" def test_iteration_loop_and_human_input_handlers(self): pipeline = _make_pipeline() @@ -602,7 +700,7 @@ def cancel(self): return None monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", + "services.workflow.execution.adapters.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", _Publisher, ) @@ -633,7 +731,7 @@ def test_init_with_end_user_sets_role_and_system_user(self): extras={}, call_depth=0, ) - workflow = _workflow() + workflow = workflow_snapshot(_workflow()) queue_manager = SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None) end_user = _end_user(end_user_id="end-user-id", session_id="session-id") @@ -644,6 +742,9 @@ def test_init_with_end_user_sets_role_and_system_user(self): user=end_user, stream=False, draft_var_saver_factory=lambda **kwargs: None, + contexts=MagicMock(), + logs=MagicMock(), + tool_providers=MagicMock(), ) assert pipeline._created_by_role == CreatorUserRole.END_USER @@ -755,7 +856,7 @@ def cancel(self): return None monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", + "services.workflow.execution.adapters.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", _Publisher, ) @@ -789,37 +890,13 @@ def cancel(self): return None monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", + "services.workflow.execution.adapters.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", _Publisher, ) with pytest.raises(RuntimeError, match="tts failure"): list(pipeline._wrapper_process_stream_response()) - @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) - def test_database_session_rolls_back_on_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_engine, sqlite_session: Session - ): - pipeline = _make_pipeline() - pipeline._application_generate_entity.invoke_from = InvokeFrom.SERVICE_API - pipeline._user_id = "user" - monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.db", - SimpleNamespace(engine=sqlite_engine), - ) - - def persist_then_fail() -> None: - with pipeline._database_session() as session: - pipeline._save_workflow_app_log(session=session, workflow_run_id="run-id") - session.flush() - raise RuntimeError("db error") - - with pytest.raises(RuntimeError, match="db error"): - persist_then_fail() - - sqlite_session.expire_all() - assert sqlite_session.scalar(select(WorkflowAppLog)) is None - def test_node_retry_and_started_handlers_cover_none_and_value(self): pipeline = _make_pipeline() pipeline._workflow_execution_id = "run-id" @@ -983,30 +1060,28 @@ def test_process_stream_response_break_paths(self): pipeline._handle_workflow_failed_and_stop_events = lambda event, **kwargs: iter(["stopped"]) assert list(pipeline._process_stream_response()) == ["stopped"] - @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) - def test_save_workflow_app_log_covers_invoke_from_variants(self, sqlite_session: Session): + def test_save_workflow_app_log_covers_invoke_from_variants(self): pipeline = _make_pipeline() pipeline._user_id = "user-id" pipeline._application_generate_entity.invoke_from = InvokeFrom.EXPLORE - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id") + pipeline._save_workflow_app_log(workflow_run_id="run-id") pipeline._application_generate_entity.invoke_from = InvokeFrom.WEB_APP - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id-2") - sqlite_session.flush() - saved_logs = sqlite_session.scalars(select(WorkflowAppLog).order_by(WorkflowAppLog.workflow_run_id)).all() - assert [log.created_from for log in saved_logs] == ["installed-app", "web-app"] + pipeline._save_workflow_app_log(workflow_run_id="run-id-2") + assert [call.kwargs["created_from"] for call in pipeline._logs.record.call_args_list] == [ + "installed-app", + "web-app", + ] - count_before = len(saved_logs) + count_before = pipeline._logs.record.call_count pipeline._application_generate_entity.invoke_from = InvokeFrom.DEBUGGER - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id-3") - sqlite_session.flush() - assert len(sqlite_session.scalars(select(WorkflowAppLog)).all()) == count_before + pipeline._save_workflow_app_log(workflow_run_id="run-id-3") + assert pipeline._logs.record.call_count == count_before pipeline._application_generate_entity.invoke_from = InvokeFrom.WEB_APP - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id=None) - sqlite_session.flush() - assert len(sqlite_session.scalars(select(WorkflowAppLog)).all()) == count_before + pipeline._save_workflow_app_log(workflow_run_id=None) + assert pipeline._logs.record.call_count == count_before def test_save_output_for_event_writes_draft_variables(self): pipeline = _make_pipeline() diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py b/api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py index 3034a328d8461e..cbecbb7248840c 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py @@ -2,13 +2,13 @@ from unittest.mock import Mock, patch -from core.app.apps.workflow.stop_aware_ready_queue import ( - StopAwareReadyQueue, - attach_stop_aware_ready_queue, -) from graphon.engine.ready_queue import StartTask from graphon.runtime import RuntimeState, VariablePool from graphon.runtime.execution import GraphExecution +from services.workflow.execution.adapters.workflow.stop_aware_ready_queue import ( + StopAwareReadyQueue, + attach_stop_aware_ready_queue, +) def _start_task(node_id: str = "next-node") -> StartTask: @@ -25,7 +25,7 @@ def test_ready_queue_accepts_work_while_run_is_active() -> None: task = _start_task() with patch( - "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", + "services.workflow.execution.adapters.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", return_value=False, ): queue.put(task) @@ -38,7 +38,7 @@ def test_ready_queue_rejects_next_node_after_graph_abort() -> None: queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution(aborted=True)) with patch( - "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", + "services.workflow.execution.adapters.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", return_value=False, ): queue.put(_start_task()) @@ -51,7 +51,7 @@ def test_ready_queue_rejects_next_node_when_stop_flag_is_set() -> None: queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution()) with patch( - "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", + "services.workflow.execution.adapters.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", return_value=True, ) as stop_flag: queue.put(_start_task("code-node")) diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py b/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py index b492ca7f62c28a..90344e4dc88713 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py @@ -96,34 +96,6 @@ def test_handle_error_preserves_agent_backend_run_failed_error(self, pipeline): assert "Knowledge retrieval failed" in str(err) assert "agent_run_id=run-1" in str(err) - @pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True) - def test_handle_error_updates_message_when_found(self, pipeline, sqlite_session: Session): - event = QueueErrorEvent(error=ValueError("oops")) - _persist_message(sqlite_session, message_id="msg-1") - - err = pipeline.handle_error(event=event, session=sqlite_session, message_id="msg-1") - - assert err is event.error - sqlite_session.flush() - sqlite_session.expire_all() - updated_message = sqlite_session.get(Message, "msg-1") - assert updated_message is not None - assert updated_message.status == MessageStatus.ERROR - assert updated_message.error == "oops" - - @pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True) - def test_handle_error_returns_err_when_message_missing(self, pipeline, sqlite_session: Session): - event = QueueErrorEvent(error=ValueError("oops")) - _persist_message(sqlite_session, message_id="other-message") - - err = pipeline.handle_error(event=event, session=sqlite_session, message_id="msg-1") - - assert err is event.error - untouched_message = sqlite_session.get(Message, "other-message") - assert untouched_message is not None - assert untouched_message.status == MessageStatus.NORMAL - assert untouched_message.error is None - def test_error_to_stream_response_and_ping(self, pipeline): error_response = pipeline.error_to_stream_response(ValueError("boom")) ping_response = pipeline.ping_stream_response() diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py index f6ecac635fe7d3..b887fb13136367 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py @@ -13,8 +13,7 @@ import pytest from sqlalchemy import Engine as _easy_ui_based_generate_task_pipeline_Engine from sqlalchemy import event -from sqlalchemy.engine import Engine as _easy_ui_message_end_files_Engine -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from core.app.app_config.entities import ( AppAdditionalFeatures, @@ -58,20 +57,22 @@ PingStreamResponse, StreamEvent, ) -from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk from core.ops.entities.trace_entity import TraceTaskName from core.ops.ops_trace_manager import TraceQueueManager -from extensions.ext_redis import RedisClientWrapper +from extensions import ext_redis from extensions.storage.storage_type import StorageType from graphon.file import FileTransferMethod, FileType from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.llm_entities import LLMResult as RuntimeLLMResult from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, TextPromptMessageContent from models.enums import ConversationFromSource, CreatorUserRole -from models.model import AppMode, Conversation, Message, MessageAgentThought, MessageFile, UploadFile +from models.model import App, AppMode, Conversation, Message, MessageAgentThought, MessageFile, UploadFile +from repositories.app.generation_repository import AppGenerationRepository +from services.app.generation.adapters.message_pipeline import EasyUIBasedGenerateTaskPipeline +from services.errors.message import MessageNotExistsError from tests.helpers.unit.core.model_fixtures import make_model_config -from tests.helpers.unit.model_factories import make_upload_file +from tests.helpers.unit.model_factories import make_app, make_conversation, make_message, make_upload_file class _QueueManager(AppQueueManager): @@ -213,6 +214,7 @@ def task_state(self): @pytest.fixture def pipeline( self, + app_records, application_generate_entity, queue_manager, conversation, @@ -222,6 +224,7 @@ def pipeline( ): """Create an EasyUIBasedGenerateTaskPipeline instance with mocked dependencies.""" pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=application_generate_entity, queue_manager=queue_manager, conversation=conversation, @@ -337,9 +340,7 @@ def test_message_end_event(self, pipeline, mock_message_cycle_manager, task_stat # Assert assert len(responses) == 1 assert task_state.llm_result == result - session = pipeline._save_message.call_args.kwargs["session"] - assert isinstance(session, Session) - assert committed_sessions == [session] + assert "session" not in pipeline._save_message.call_args.kwargs pipeline._message_end_to_stream_response.assert_called_once() def test_error_event(self, pipeline, committed_sessions): @@ -359,12 +360,8 @@ def test_error_event(self, pipeline, committed_sessions): # Assert assert len(responses) == 1 - session = pipeline.handle_error.call_args.kwargs["session"] - assert isinstance(session, Session) - assert committed_sessions == [session] - pipeline.handle_error.assert_called_once_with( - event=error_event, session=session, message_id="test-message-id" - ) + assert len(committed_sessions) == 1 + pipeline.handle_error.assert_called_once_with(event=error_event) pipeline.error_to_stream_response.assert_called_once() def test_ping_event(self, pipeline): @@ -448,10 +445,8 @@ def test_trace_manager_passed_to_save_message(self, pipeline, committed_sessions list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) # Assert - session = pipeline._save_message.call_args.kwargs["session"] - assert isinstance(session, Session) - assert committed_sessions == [session] - pipeline._save_message.assert_called_once_with(session=session, trace_manager=trace_manager) + assert "session" not in pipeline._save_message.call_args.kwargs + pipeline._save_message.assert_called_once_with(trace_manager=trace_manager, preserve_existing_usage=False) def test_multiple_events_sequence(self, pipeline, mock_message_cycle_manager, task_state): """Test handling multiple events in sequence.""" @@ -496,6 +491,7 @@ def test_multiple_events_sequence(self, pipeline, mock_message_cycle_manager, ta class _DummyModelConf: def __init__(self) -> None: self.model = "mock" + self.mode = "chat" class _UnknownQueueEvent: @@ -713,6 +709,7 @@ def _make_entity(entity_cls, app_mode: AppMode): def _make_pipeline( + app_records, entity_cls: type[ChatAppGenerateEntity] | type[CompletionAppGenerateEntity] = ChatAppGenerateEntity, app_mode: AppMode = AppMode.CHAT, *, @@ -720,6 +717,7 @@ def _make_pipeline( ) -> tuple[EasyUIBasedGenerateTaskPipeline, _FakeQueueManager]: queue_manager = _FakeQueueManager() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(entity_cls, app_mode), queue_manager=queue_manager, conversation=_make_conversation(app_mode), @@ -742,11 +740,12 @@ def _set_method(obj: object, name: str, value: object) -> None: class TestEasyUiBasedGenerateTaskPipelineCoreCases: class TestEasyUiBasedGenerateTaskPipeline: - def test_to_blocking_response_chat(self): + def test_to_blocking_response_chat(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -762,11 +761,12 @@ def _gen(): assert response.data.answer == "answer" - def test_to_blocking_response_completion(self): + def test_to_blocking_response_completion(self, app_records): conversation = _make_conversation(AppMode.COMPLETION) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -782,11 +782,12 @@ def _gen(): assert response.data.answer == "answer" - def test_listen_audio_msg_returns_none_when_no_publisher(self): + def test_listen_audio_msg_returns_none_when_no_publisher(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -796,11 +797,12 @@ def test_listen_audio_msg_returns_none_when_no_publisher(self): assert pipeline._listen_audio_msg(publisher=None, task_id="task") is None - def test_process_stream_response_handles_chunks_and_end(self, monkeypatch: pytest.MonkeyPatch): + def test_process_stream_response_handles_chunks_and_end(self, app_records, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -845,24 +847,6 @@ def _message_end() -> MessageEndStreamResponse: _set_method(pipeline, "_message_end_to_stream_response", _message_end) _set_method(pipeline, "_save_message", lambda **kwargs: None) - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - responses = list(pipeline._process_stream_response(publisher=None)) message_response = next(item for item in responses if isinstance(item, MessageStreamResponse)) @@ -873,11 +857,12 @@ def commit(self): assert isinstance(responses[-1], MessageEndStreamResponse) assert pipeline._task_state.llm_result.message.content == "done" - def test_handle_output_moderation_chunk_directs_output(self): + def test_handle_output_moderation_chunk_directs_output(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -901,7 +886,7 @@ def get_final_output(self): assert any(isinstance(event, QueueLLMChunkEvent) for event in events) assert any(isinstance(event, QueueStopEvent) for event in events) - def test_handle_stop_updates_usage(self, monkeypatch: pytest.MonkeyPatch): + def test_handle_stop_updates_usage(self, app_records, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) message = _make_message() @@ -940,6 +925,7 @@ def __init__(self) -> None: ) pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=application_generate_entity, queue_manager=_FakeQueueManager(), conversation=conversation, @@ -960,7 +946,7 @@ def get_llm_num_tokens(self, messages): return 10 if len(calls) == 1 else 5 monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.ModelInstance", + "services.app.generation.adapters.message_pipeline.ModelInstance", _FakeModelInstance, ) @@ -969,11 +955,12 @@ def get_llm_num_tokens(self, messages): assert pipeline._task_state.llm_result.usage.prompt_tokens == 10 assert pipeline._task_state.llm_result.usage.completion_tokens == 5 - def test_record_files_builds_file_payloads(self, monkeypatch: pytest.MonkeyPatch): + def test_record_files_builds_file_payloads(self, app_records, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1012,31 +999,6 @@ def test_record_files_builds_file_payloads(self, monkeypatch: pytest.MonkeyPatch ) ] - class _Result: - def __init__(self, items): - self._items = items - - def all(self): - return self._items - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - self.calls = 0 - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - self.calls += 1 - return _Result(message_files if self.calls == 1 else upload_files) - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) monkeypatch.setattr( "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", lambda **kwargs: "signed-url", @@ -1046,17 +1008,23 @@ def scalars(self, *args, **kwargs): lambda **kwargs: "signed-tool", ) + monkeypatch.setattr( + app_records, "message_files", lambda identity: (message_files, {file.id: file for file in upload_files}) + ) response = pipeline._message_end_to_stream_response() files = response.files assert files assert len(files) == 3 - def test_process_stream_response_handles_annotation_and_error(self, monkeypatch: pytest.MonkeyPatch): + def test_process_stream_response_handles_annotation_and_error( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1097,7 +1065,7 @@ def _agent_message_response(answer: str, message_id: str) -> AgentMessageStreamR _set_method( pipeline._message_cycle_manager, "handle_annotation_reply", - lambda event, session: _AnnotationReply(content="annotated"), + lambda event: _AnnotationReply(content="annotated"), ) _set_method(pipeline, "_agent_thought_to_stream_response", _agent_thought_response) _set_method(pipeline._message_cycle_manager, "message_file_to_stream_response", _file_response) @@ -1105,24 +1073,6 @@ def _agent_message_response(answer: str, message_id: str) -> AgentMessageStreamR _set_method(pipeline, "handle_error", lambda **kwargs: ValueError("boom")) _set_method(pipeline, "error_to_stream_response", lambda err: ErrorStreamResponse(task_id="task", err=err)) - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - responses = list(pipeline._process_stream_response(publisher=None)) assert any(isinstance(response, AgentThoughtStreamResponse) for response in responses) @@ -1135,13 +1085,16 @@ def commit(self): assert isinstance(responses[-1].err, ValueError) assert pipeline._task_state.llm_result.message.content == "annotated" - def test_process_stream_response_error_event_adds_trace_task(self, monkeypatch: pytest.MonkeyPatch): + def test_process_stream_response_error_event_adds_trace_task( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) application_generate_entity.extras = {"trace_session_id": "session-1"} pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=application_generate_entity, queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1155,21 +1108,6 @@ def test_process_stream_response_error_event_adds_trace_task(self, monkeypatch: trace_manager_double = _TraceManagerDouble() trace_manager = cast(TraceQueueManager, trace_manager_double) - class _Session: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) assert len(responses) == 1 @@ -1181,11 +1119,12 @@ def commit(self): assert trace_task.message_id == "msg" assert trace_task.kwargs["trace_session_id"] == "session-1" - def test_agent_thought_to_stream_response_returns_payload(self, monkeypatch: pytest.MonkeyPatch): + def test_agent_thought_to_stream_response_returns_payload(self, app_records, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1195,35 +1134,21 @@ def test_agent_thought_to_stream_response_returns_payload(self, monkeypatch: pyt agent_thought = _agent_thought() - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalar(self, *args, **kwargs): - return agent_thought - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - + monkeypatch.setattr(app_records, "agent_thought", lambda identity, thought_id: agent_thought) response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="thought")) assert response is not None assert response.id == "thought" assert response.thought == "t" - def test_agent_thought_to_stream_response_normalizes_null_display_fields(self, monkeypatch: pytest.MonkeyPatch): + def test_agent_thought_to_stream_response_normalizes_null_display_fields( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1239,23 +1164,7 @@ def test_agent_thought_to_stream_response_normalizes_null_display_fields(self, m agent_thought.tool_input = None agent_thought.message_files = None - class _Session: - def __init__(self, *args, **kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalar(self, *args, **kwargs): - return agent_thought - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) + monkeypatch.setattr(app_records, "agent_thought", lambda identity, thought_id: agent_thought) response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="thought")) assert response is not None @@ -1265,10 +1174,11 @@ def scalar(self, *args, **kwargs): assert response.tool_input == "" assert response.model_dump(mode="json")["message_files"] == [] - def test_process_routes_to_stream_and_starts_conversation_name_generation(self): + def test_process_routes_to_stream_and_starts_conversation_name_generation(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1290,10 +1200,11 @@ def test_process_routes_to_stream_and_starts_conversation_name_generation(self): conversation_id="conv", query="hello", message_id="msg" ) - def test_process_routes_to_blocking_for_completion_mode(self): + def test_process_routes_to_blocking_for_completion_mode(self, app_records): conversation = _make_conversation(AppMode.COMPLETION) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1313,10 +1224,11 @@ def test_process_routes_to_blocking_for_completion_mode(self): assert result == "blocking" pipeline._message_cycle_manager.generate_conversation_name.assert_not_called() - def test_to_blocking_response_raises_error_stream_exception(self): + def test_to_blocking_response_raises_error_stream_exception(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1330,10 +1242,11 @@ def _gen(): with pytest.raises(ValueError, match="stream error"): pipeline._to_blocking_response(_gen()) - def test_to_blocking_response_raises_when_generator_ends_without_message_end(self): + def test_to_blocking_response_raises_when_generator_ends_without_message_end(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1347,10 +1260,11 @@ def _gen(): with pytest.raises(RuntimeError, match="queue listening stopped unexpectedly"): pipeline._to_blocking_response(_gen()) - def test_to_stream_response_wraps_completion_stream_events(self): + def test_to_stream_response_wraps_completion_stream_events(self, app_records): conversation = _make_conversation(AppMode.COMPLETION) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1366,10 +1280,11 @@ def _gen(): assert isinstance(response, CompletionAppStreamResponse) assert response.message_id == "msg" - def test_to_stream_response_wraps_chat_stream_events(self): + def test_to_stream_response_wraps_chat_stream_events(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1385,10 +1300,11 @@ def _gen(): assert isinstance(response, ChatbotAppStreamResponse) assert response.conversation_id == "conv" - def test_listen_audio_msg_returns_audio_response_for_non_finish_audio(self): + def test_listen_audio_msg_returns_audio_response_for_non_finish_audio(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1405,10 +1321,11 @@ def test_listen_audio_msg_returns_audio_response_for_non_finish_audio(self): assert isinstance(response, MessageAudioStreamResponse) assert response.audio == "abc" - def test_listen_audio_msg_returns_none_for_finish_audio(self): + def test_listen_audio_msg_returns_none_for_finish_audio(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1422,10 +1339,11 @@ def test_listen_audio_msg_returns_none_for_finish_audio(self): assert pipeline._listen_audio_msg(publisher=publisher, task_id="task") is None - def test_wrapper_process_stream_response_without_tts_publisher(self): + def test_wrapper_process_stream_response_without_tts_publisher(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1439,7 +1357,7 @@ def test_wrapper_process_stream_response_without_tts_publisher(self): assert responses == [payload] - def test_wrapper_process_stream_response_with_tts_publisher(self, monkeypatch: pytest.MonkeyPatch): + def test_wrapper_process_stream_response_with_tts_publisher(self, app_records, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) message = _make_message() entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) @@ -1447,6 +1365,7 @@ def test_wrapper_process_stream_response_with_tts_publisher(self, monkeypatch: p "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} } pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=entity, queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1471,7 +1390,7 @@ def cancel(self): _set_method(pipeline, "_listen_audio_msg", lambda publisher, task_id: next(audio_calls)) _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([payload])) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + "services.app.generation.adapters.message_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: _Publisher(), ) @@ -1482,7 +1401,9 @@ def cancel(self): assert isinstance(responses[-1], MessageAudioEndStreamResponse) assert "audio_type" not in responses[-1].model_dump() - def test_wrapper_process_stream_response_waits_for_the_audio_terminal(self, monkeypatch: pytest.MonkeyPatch): + def test_wrapper_process_stream_response_waits_for_the_audio_terminal( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) @@ -1490,6 +1411,7 @@ def test_wrapper_process_stream_response_waits_for_the_audio_terminal(self, monk "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} } pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=entity, queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1513,7 +1435,7 @@ def cancel(self): _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + "services.app.generation.adapters.message_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: _Publisher(), ) responses = list(pipeline._wrapper_process_stream_response()) @@ -1522,7 +1444,7 @@ def cancel(self): assert isinstance(responses[-1], MessageAudioEndStreamResponse) def test_wrapper_process_stream_response_ends_audio_before_reporting_tts_error( - self, monkeypatch: pytest.MonkeyPatch + self, app_records, monkeypatch: pytest.MonkeyPatch ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() @@ -1531,6 +1453,7 @@ def test_wrapper_process_stream_response_ends_audio_before_reporting_tts_error( "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} } pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=entity, queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1552,7 +1475,7 @@ def cancel(self): _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + "services.app.generation.adapters.message_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: _Publisher(), ) @@ -1562,7 +1485,9 @@ def cancel(self): assert isinstance(responses[-1], ErrorStreamResponse) assert responses[-1].err is error - def test_wrapper_process_stream_response_ends_audio_before_a_main_error(self, monkeypatch: pytest.MonkeyPatch): + def test_wrapper_process_stream_response_ends_audio_before_a_main_error( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) @@ -1570,6 +1495,7 @@ def test_wrapper_process_stream_response_ends_audio_before_a_main_error(self, mo "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} } pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=entity, queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1581,7 +1507,7 @@ def test_wrapper_process_stream_response_ends_audio_before_a_main_error(self, mo error = ErrorStreamResponse(task_id="task", err=RuntimeError("generation failed")) _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([error])) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + "services.app.generation.adapters.message_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: publisher, ) @@ -1593,11 +1519,12 @@ def test_wrapper_process_stream_response_ends_audio_before_a_main_error(self, mo publisher.publish.assert_not_called() def test_process_stream_response_handles_stop_event_and_output_replacement( - self, monkeypatch: pytest.MonkeyPatch + self, app_records, monkeypatch: pytest.MonkeyPatch ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1620,24 +1547,6 @@ def test_process_stream_response_handles_stop_event_and_output_replacement( lambda: MessageEndStreamResponse(task_id="task", id="msg"), ) - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - responses = list(pipeline._process_stream_response(publisher=None)) assert isinstance(responses[0], MessageReplaceStreamResponse) @@ -1645,10 +1554,11 @@ def commit(self): assert isinstance(responses[1], MessageEndStreamResponse) pipeline._handle_stop.assert_called_once() - def test_process_stream_response_handles_retriever_unknown_and_empty_chunk(self): + def test_process_stream_response_handles_retriever_unknown_and_empty_chunk(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1681,10 +1591,11 @@ def _handle_retriever_resources(event): assert responses == [] assert handled["retriever"] == 1 - def test_process_stream_response_skips_when_output_moderation_directs_chunk(self): + def test_process_stream_response_skips_when_output_moderation_directs_chunk(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1703,10 +1614,11 @@ def test_process_stream_response_skips_when_output_moderation_directs_chunk(self assert responses == [] - def test_process_stream_response_ignores_unsupported_chunk_content_types(self): + def test_process_stream_response_ignores_unsupported_chunk_content_types(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1731,10 +1643,11 @@ def test_process_stream_response_ignores_unsupported_chunk_content_types(self): assert responses[0].answer == "ok" assert pipeline._task_state.llm_result.message.content == "ok" - def test_process_stream_response_skips_none_chunk_content(self): + def test_process_stream_response_skips_none_chunk_content(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1755,10 +1668,11 @@ def test_process_stream_response_skips_none_chunk_content(self): pipeline._message_cycle_manager.message_to_stream_response.assert_not_called() assert pipeline._task_state.llm_result.message.content == "" - def test_process_stream_response_reaches_post_loop_branch_with_thread_reference(self): + def test_process_stream_response_reaches_post_loop_branch_with_thread_reference(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1771,13 +1685,14 @@ def test_process_stream_response_reaches_post_loop_branch_with_thread_reference( assert list(pipeline._process_stream_response(publisher=None)) == [] def test_save_message_persists_fields_and_emits_trace( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + self, app_records, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) application_generate_entity.extras = {"trace_session_id": "session-1"} pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=application_generate_entity, queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1796,34 +1711,33 @@ def test_save_message_persists_fields_and_emits_trace( conversation_obj = _make_conversation(AppMode.CHAT) session = sqlite_session session.add_all([conversation_obj, message_obj]) - session.flush() + session.commit() trace_manager_double = _TraceManagerDouble() trace_manager = cast(TraceQueueManager, trace_manager_double) sent_payloads: list[tuple[tuple[object, ...], dict[str, object]]] = [] monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", - lambda mode, prompt_messages: "serialized-prompt", + "services.app.generation.adapters.message_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", + lambda mode, prompt_messages: [{"role": "user", "text": "serialized-prompt"}], ) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptTemplateParser.remove_template_variables", + "services.app.generation.adapters.message_pipeline.PromptTemplateParser.remove_template_variables", lambda text: text.replace("{{name}}", "").strip(), ) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.naive_utc_now", + "repositories.app.generation_repository.naive_utc_now", lambda: datetime(2024, 1, 1, tzinfo=UTC), ) + monkeypatch.setattr("services.app.generation.adapters.message_pipeline.time.perf_counter", lambda: 15.0) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.time.perf_counter", lambda: 15.0 - ) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", + "services.app.generation.adapters.message_pipeline.message_was_created.send", lambda *args, **kwargs: sent_payloads.append((args, kwargs)), ) - pipeline._save_message(session=session, trace_manager=trace_manager) + pipeline._save_message(trace_manager=trace_manager) - assert message_obj.message == "serialized-prompt" + session.refresh(message_obj) + assert message_obj.message == [{"role": "user", "text": "serialized-prompt"}] assert message_obj.answer == "hello" assert message_obj.provider_response_latency == 5.0 trace_manager_double.add_trace_task.assert_called_once() @@ -1835,9 +1749,10 @@ def test_save_message_persists_fields_and_emits_trace( assert len(sent_payloads) == 1 def test_save_stopped_message_preserves_backend_reported_usage( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + self, app_records, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=_make_conversation(AppMode.CHAT), @@ -1853,18 +1768,19 @@ def test_save_stopped_message_preserves_backend_reported_usage( message.total_price = Decimal("0.21") message.message_metadata = json.dumps({"usage": {"total_tokens": 21}}) sqlite_session.add_all([_make_conversation(AppMode.CHAT), message]) - sqlite_session.flush() + sqlite_session.commit() monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", - lambda _mode, _prompt_messages: "", + "services.app.generation.adapters.message_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", + lambda _mode, _prompt_messages: [], ) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", + "services.app.generation.adapters.message_pipeline.message_was_created.send", lambda *_args, **_kwargs: None, ) - pipeline._save_message(session=sqlite_session, preserve_existing_usage=True) + pipeline._save_message(preserve_existing_usage=True) + sqlite_session.refresh(message) assert message.message_tokens == 13 assert message.answer_tokens == 8 assert message.total_price == Decimal("0.21") @@ -1874,11 +1790,13 @@ def test_save_stopped_message_preserves_backend_reported_usage( @pytest.mark.parametrize("metadata", ["{", "[]"]) def test_save_stopped_message_recovers_invalid_existing_metadata( self, + app_records, metadata: str, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, ) -> None: pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=_make_conversation(AppMode.CHAT), @@ -1894,41 +1812,42 @@ def test_save_stopped_message_recovers_invalid_existing_metadata( message.total_price = Decimal("0.21") message.message_metadata = metadata sqlite_session.add_all([_make_conversation(AppMode.CHAT), message]) - sqlite_session.flush() + sqlite_session.commit() monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", - lambda _mode, _prompt_messages: "", + "services.app.generation.adapters.message_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", + lambda _mode, _prompt_messages: [], ) monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", + "services.app.generation.adapters.message_pipeline.message_was_created.send", lambda *_args, **_kwargs: None, ) - pipeline._save_message(session=sqlite_session, preserve_existing_usage=True) + pipeline._save_message(preserve_existing_usage=True) + sqlite_session.refresh(message) assert message.message_tokens == 13 assert message.answer_tokens == 8 assert message.total_price == Decimal("0.21") - def test_save_message_raises_when_message_not_found(self, sqlite_session: Session): + def test_save_message_raises_when_message_not_found(self, app_records, sqlite_session: Session): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, message=message, stream=False, ) - session = sqlite_session + with pytest.raises(MessageNotExistsError): + pipeline._save_message() - with pytest.raises(ValueError, match="message msg not found"): - pipeline._save_message(session=session) - - def test_save_message_raises_when_conversation_not_found(self, sqlite_session: Session): + def test_save_message_raises_when_conversation_not_found(self, app_records, sqlite_session: Session): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1937,15 +1856,18 @@ def test_save_message_raises_when_conversation_not_found(self, sqlite_session: S ) session = sqlite_session session.add(_make_message()) - session.flush() + session.commit() - with pytest.raises(ValueError, match="Conversation conv not found"): - pipeline._save_message(session=session) + with pytest.raises(MessageNotExistsError): + pipeline._save_message() - def test_message_end_to_stream_response_includes_usage_metadata(self, monkeypatch: pytest.MonkeyPatch): + def test_message_end_to_stream_response_includes_usage_metadata( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1954,38 +1876,19 @@ def test_message_end_to_stream_response_includes_usage_metadata(self, monkeypatc ) pipeline._task_state.llm_result.usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 2}) - class _Result: - def all(self): - return [] - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - return _Result() - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - response = pipeline._message_end_to_stream_response() assert response.id == "msg" usage_metadata = cast(dict[str, object], response.metadata["usage"]) assert usage_metadata["prompt_tokens"] == 1 - def test_record_files_returns_empty_list_when_message_has_no_files(self, monkeypatch: pytest.MonkeyPatch): + def test_record_files_returns_empty_list_when_message_has_no_files( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -1993,36 +1896,17 @@ def test_record_files_returns_empty_list_when_message_has_no_files(self, monkeyp stream=False, ) - class _Result: - def all(self): - return [] - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - return _Result() - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - response = pipeline._message_end_to_stream_response() assert response.files == [] - def test_record_files_handles_local_fallback_and_tool_url_variants(self, monkeypatch: pytest.MonkeyPatch): + def test_record_files_handles_local_fallback_and_tool_url_variants( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -2053,31 +1937,6 @@ def test_record_files_handles_local_fallback_and_tool_url_variants(self, monkeyp ), ] - class _Result: - def __init__(self, items): - self._items = items - - def all(self): - return self._items - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - self.calls = 0 - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - self.calls += 1 - return _Result(message_files if self.calls == 1 else []) - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) monkeypatch.setattr( "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", lambda **kwargs: "local-fallback-signed", @@ -2087,6 +1946,7 @@ def scalars(self, *args, **kwargs): lambda **kwargs: "tool-signed", ) + monkeypatch.setattr(app_records, "message_files", lambda identity: (message_files, {})) response = pipeline._message_end_to_stream_response() files = response.files @@ -2095,10 +1955,11 @@ def scalars(self, *args, **kwargs): assert files[1]["filename"] == "file.txt" assert files[2]["extension"] == ".bin" - def test_agent_message_to_stream_response_builds_payload(self): + def test_agent_message_to_stream_response_builds_payload(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -2111,10 +1972,13 @@ def test_agent_message_to_stream_response_builds_payload(self): assert response.id == "msg" assert response.answer == "hello" - def test_agent_thought_to_stream_response_returns_none_when_not_found(self, monkeypatch: pytest.MonkeyPatch): + def test_agent_thought_to_stream_response_returns_none_when_not_found( + self, app_records, monkeypatch: pytest.MonkeyPatch + ): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -2122,32 +1986,15 @@ def test_agent_thought_to_stream_response_returns_none_when_not_found(self, monk stream=True, ) - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalar(self, *args, **kwargs): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="missing")) assert response is None - def test_handle_output_moderation_chunk_appends_token_when_not_directing(self): + def test_handle_output_moderation_chunk_appends_token_when_not_directing(self, app_records): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( + records=app_records, application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), queue_manager=_FakeQueueManager(), conversation=conversation, @@ -2170,6 +2017,12 @@ def append_new_token(self, text): assert result is False assert appended_tokens == ["next-token"] + @staticmethod + @pytest.fixture(autouse=True) + def generation_app(sqlite_session): + sqlite_session.add(make_app(app_id="app", tenant_id="tenant")) + sqlite_session.commit() + """Exercise message-end file serialization against persisted SQLite rows. @@ -2178,7 +2031,7 @@ def append_new_token(self, text): """ -SQLITE_MODELS = (MessageFile, UploadFile) +SQLITE_MODELS = (App, Conversation, Message, MessageFile, UploadFile) class TestEasyUiMessageEndFilesCases: @@ -2190,27 +2043,39 @@ class TestEasyUiMessageEndFilesCases: class TestMessageEndStreamResponseFiles: """Verify message-end file payloads from actual ORM query results.""" - @pytest.fixture(autouse=True) - def bind_sqlite_engine( - self, sqlite_engine: _easy_ui_message_end_files_Engine, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Bind sessions opened by the pipeline to the per-test SQLite engine.""" - - sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - monkeypatch.setattr("core.db.session_factory._session_maker", sqlite_session_maker) - @pytest.fixture def pipeline( - self, monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock] + self, + monkeypatch: pytest.MonkeyPatch, + redis_transport: tuple[ext_redis.RedisClientWrapper, MagicMock], + app_records: AppGenerationRepository, + sqlite_session: Session, ) -> EasyUIBasedGenerateTaskPipeline: - """Construct the pipeline with validated request, model, and task state.""" + """Construct the service adapter with real request models and injected records.""" monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0]) + app = make_app() + conversation = make_conversation(inputs={}, from_source=ConversationFromSource.API) + message = make_message( + message_id=str(uuid.uuid4()), + inputs={}, + query="hello", + message={}, + answer="", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + created_at=datetime.now(), + ) + sqlite_session.add_all([app, conversation, message]) + sqlite_session.commit() + request = ChatAppGenerateEntity( task_id=str(uuid.uuid4()), app_config=EasyUIBasedAppConfig( - tenant_id="tenant-id", - app_id="app-id", - app_mode=AppMode.CHAT, + tenant_id=app.tenant_id, + app_id=app.id, + app_mode=app.mode, app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, app_model_config_dict={}, model=ModelConfigEntity(provider="test-provider", model="test-model"), @@ -2223,22 +2088,12 @@ def pipeline( stream=True, invoke_from=InvokeFrom.WEB_APP, ) - conversation = Conversation( - id=str(uuid.uuid4()), - app_id="app-id", - mode=AppMode.CHAT, - name="Conversation", - from_source=ConversationFromSource.API, - _inputs={}, - ) - message = Message(id=str(uuid.uuid4())) - message.created_at = datetime.now() queue = MessageBasedAppQueueManager( task_id=request.task_id, user_id=request.user_id, invoke_from=request.invoke_from, conversation_id=conversation.id, - app_mode=AppMode.CHAT, + app_mode=app.mode, message_id=message.id, ) return EasyUIBasedGenerateTaskPipeline( @@ -2247,6 +2102,7 @@ def pipeline( conversation=conversation, message=message, stream=True, + records=app_records, ) @staticmethod @@ -2298,7 +2154,7 @@ def upload_file(self, message_file_local: MessageFile) -> UploadFile: """Create upload metadata matching the local message-file reference.""" upload = UploadFile( - tenant_id=str(uuid.uuid4()), + tenant_id="tenant-1", storage_type=StorageType.LOCAL, key="uploads/test_image.png", name="test_image.png", diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py index f50355ace268ab..52b95f557c1fe7 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py @@ -1,7 +1,6 @@ """Unit tests for the message cycle manager optimization.""" import logging -from collections.abc import Iterator from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import Mock, patch @@ -9,21 +8,18 @@ import pytest from flask import Flask, current_app from sqlalchemy import Engine, event, select -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from core.app.app_config.entities import AppAdditionalFeatures from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueRetrieverResourcesEvent from core.app.entities.task_entities import MessageStreamResponse, StreamEvent, TaskStateMetadata, WorkflowTaskState -from core.app.task_pipeline import message_cycle_manager as message_cycle_manager_module -from core.app.task_pipeline.message_cycle_manager import MessageCycleManager from core.rag.entities import RetrievalSourceMetadata from graphon.file import FileTransferMethod, FileType -from models import model as model_module from models.account import Account -from models.base import TypeBase from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo -from models.model import App, Conversation, MessageAnnotation, MessageFile -from tests.helpers.unit.model_factories import make_app, make_conversation +from models.model import App, Conversation, Message, MessageAnnotation, MessageFile +from services.app.generation.adapters.message_cycle import MessageCycleManager +from tests.helpers.unit.model_factories import make_app, make_conversation, make_message @dataclass(frozen=True) @@ -31,26 +27,9 @@ class _AppConfigWithFeatures: additional_features: AppAdditionalFeatures -@dataclass(frozen=True) -class _SQLiteDb: - engine: Engine - session: Session - - @pytest.fixture -def cycle_db(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: - """Bind request-owned and cycle-manager-owned sessions to isolated SQLite.""" - TypeBase.metadata.create_all( - sqlite_engine, - tables=[App.__table__, Conversation.__table__, MessageFile.__table__], - ) - owned_session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - with owned_session_factory() as request_session: - sqlite_db = _SQLiteDb(engine=sqlite_engine, session=request_session) - monkeypatch.setattr(message_cycle_manager_module, "db", sqlite_db) - monkeypatch.setattr(model_module, "db", sqlite_db) - monkeypatch.setattr(message_cycle_manager_module.session_factory, "create_session", owned_session_factory) - yield request_session +def cycle_db(sqlite_session: Session) -> Session: + return sqlite_session def _app(*, app_id: str = "app-id", tenant_id: str = "tenant-1") -> App: @@ -69,6 +48,7 @@ def _conversation(*, conversation_id: str = "conv-1", app_id: str = "app-id") -> def _message_file( + session: Session, *, file_id: str = "file-1", message_id: str = "test-message-id", @@ -76,6 +56,23 @@ def _message_file( url: str | None = "http://example.com/image.png", file_type: FileType = FileType.IMAGE, ) -> MessageFile: + if session.get(App, "app-id") is None: + session.add(_app()) + if session.get(Message, message_id) is None: + session.add( + make_message( + message_id=message_id, + app_id="app-id", + inputs={}, + query="", + message={}, + answer="", + message_unit_price=0, + answer_unit_price=0, + currency="USD", + from_source=ConversationFromSource.API, + ) + ) message_file = MessageFile( message_id=message_id, type=file_type, @@ -97,13 +94,17 @@ def mock_application_generate_entity(self): """Create a mock application generate entity.""" entity = Mock() entity.task_id = "test-task-id" + entity.app_config.tenant_id = "tenant-1" + entity.app_config.app_id = "app-id" return entity @pytest.fixture - def message_cycle_manager(self, mock_application_generate_entity): + def message_cycle_manager(self, app_records, mock_application_generate_entity): """Create a message cycle manager instance.""" task_state = Mock() - return MessageCycleManager(application_generate_entity=mock_application_generate_entity, task_state=task_state) + return MessageCycleManager( + records=app_records, application_generate_entity=mock_application_generate_entity, task_state=task_state + ) def test_get_message_event_type_with_assistant_file(self, message_cycle_manager, cycle_db: Session): """Test get_message_event_type returns MESSAGE_FILE when message has assistant-generated files. @@ -111,7 +112,11 @@ def test_get_message_event_type_with_assistant_file(self, message_cycle_manager, This ensures that AI-generated images (belongs_to='assistant') trigger the MESSAGE_FILE event, allowing the frontend to properly display generated image files with url field. """ - cycle_db.add(_message_file()) + cycle_db.add( + _message_file( + cycle_db, + ) + ) cycle_db.commit() with current_app.app_context(): @@ -128,7 +133,7 @@ def test_get_message_event_type_with_user_file(self, message_cycle_manager, cycl resulting in broken images in the chat UI. The query filters for belongs_to='assistant', so when only user files exist, the database query returns None, resulting in MESSAGE event type. """ - cycle_db.add(_message_file(belongs_to=MessageFileBelongsTo.USER)) + cycle_db.add(_message_file(cycle_db, belongs_to=MessageFileBelongsTo.USER)) cycle_db.commit() with current_app.app_context(): @@ -167,7 +172,11 @@ def record_statement(_conn, _cursor, statement, _parameters, _context, _executem def test_message_to_stream_response_with_precomputed_event_type(self, message_cycle_manager, cycle_db: Session): """MessageCycleManager.message_to_stream_response expects a valid event_type; callers should precompute it.""" - cycle_db.add(_message_file()) + cycle_db.add( + _message_file( + cycle_db, + ) + ) cycle_db.commit() with current_app.app_context(): @@ -257,7 +266,7 @@ def test_generate_conversation_name_returns_none_for_completion(self, message_cy class DummyCompletion: pass - with patch("core.app.task_pipeline.message_cycle_manager.CompletionAppGenerateEntity", DummyCompletion): + with patch("services.app.generation.adapters.message_cycle.CompletionAppGenerateEntity", DummyCompletion): message_cycle_manager._application_generate_entity = DummyCompletion() result = message_cycle_manager.generate_conversation_name(conversation_id="c1", query="hi") @@ -283,10 +292,10 @@ def start(self): with ( patch( - "core.app.task_pipeline.message_cycle_manager.current_app", + "services.app.generation.adapters.message_cycle.current_app", new=SimpleNamespace(_get_current_object=lambda: flask_app), ), - patch("core.app.task_pipeline.message_cycle_manager.Timer", DummyTimer), + patch("services.app.generation.adapters.message_cycle.Timer", DummyTimer), ): thread = message_cycle_manager.generate_conversation_name( conversation_id="conv-1", query="hello", message_id="message-1" @@ -308,7 +317,7 @@ def test_generate_conversation_name_skips_thread_when_auto_generate_disabled(sel message_cycle_manager._application_generate_entity.is_new_conversation = True message_cycle_manager._application_generate_entity.extras = {"auto_generate_conversation_name": False} - with patch("core.app.task_pipeline.message_cycle_manager.Timer") as mock_timer: + with patch("services.app.generation.adapters.message_cycle.Timer") as mock_timer: result = message_cycle_manager.generate_conversation_name(conversation_id="conv-2", query="hello") assert result is None @@ -347,8 +356,8 @@ def test_generate_conversation_name_worker_uses_cached_name( cycle_db.commit() with ( - patch("core.app.task_pipeline.message_cycle_manager.redis_client") as mock_redis, - patch("core.app.task_pipeline.message_cycle_manager.LLMGenerator") as mock_llm_generator, + patch("services.app.generation.adapters.message_cycle.redis_client") as mock_redis, + patch("services.app.generation.adapters.message_cycle.LLMGenerator") as mock_llm_generator, ): mock_redis.get.return_value = b"cached-title" @@ -371,8 +380,8 @@ def test_generate_conversation_name_worker_generates_and_caches_name( cycle_db.commit() with ( - patch("core.app.task_pipeline.message_cycle_manager.redis_client") as mock_redis, - patch("core.app.task_pipeline.message_cycle_manager.LLMGenerator") as mock_llm_generator, + patch("services.app.generation.adapters.message_cycle.redis_client") as mock_redis, + patch("services.app.generation.adapters.message_cycle.LLMGenerator") as mock_llm_generator, ): mock_redis.get.return_value = None mock_llm_generator.generate_conversation_name.return_value = "generated-title" @@ -407,12 +416,12 @@ def test_generate_conversation_name_worker_falls_back_when_generation_fails( long_query = "q" * 60 with ( - patch("core.app.task_pipeline.message_cycle_manager.redis_client") as mock_redis, - patch("core.app.task_pipeline.message_cycle_manager.LLMGenerator") as mock_llm_generator, + patch("services.app.generation.adapters.message_cycle.redis_client") as mock_redis, + patch("services.app.generation.adapters.message_cycle.LLMGenerator") as mock_llm_generator, ): mock_redis.get.return_value = None mock_llm_generator.generate_conversation_name.side_effect = RuntimeError("generation failed") - with caplog.at_level(logging.ERROR, logger="core.app.task_pipeline.message_cycle_manager"): + with caplog.at_level(logging.ERROR, logger="services.app.generation.adapters.message_cycle"): message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", long_query) assert cycle_db.in_transaction() is False @@ -435,6 +444,7 @@ def test_handle_annotation_reply_reads_author_in_one_query( account_id="acct-1", ) annotation.id = "ann-1" + sqlite_session.add(_app()) sqlite_session.add(annotation) if author_name is not None: account = Account(name=author_name, email="author@example.com") @@ -450,7 +460,7 @@ def capture_sql(_connection, _cursor, statement: str, _parameters, _context, _ex event.listen(sqlite_engine, "before_cursor_execute", capture_sql) try: result = message_cycle_manager.handle_annotation_reply( - QueueAnnotationReplyEvent(message_annotation_id="ann-1"), sqlite_session + QueueAnnotationReplyEvent(message_annotation_id="ann-1") ) finally: event.remove(sqlite_engine, "before_cursor_execute", capture_sql) @@ -469,7 +479,7 @@ def test_handle_annotation_reply_returns_none_when_missing(self, message_cycle_m """Return None and keep metadata unchanged when annotation is not found.""" message_cycle_manager._task_state = WorkflowTaskState() result = message_cycle_manager.handle_annotation_reply( - QueueAnnotationReplyEvent(message_annotation_id="missing"), sqlite_session + QueueAnnotationReplyEvent(message_annotation_id="missing") ) assert result is None @@ -508,6 +518,7 @@ def test_message_file_to_stream_response_builds_signed_url(self, message_cycle_m message_cycle_manager._application_generate_entity.task_id = "task-1" cycle_db.add( _message_file( + cycle_db, file_id="file-1", message_id="msg-1", belongs_to=None, @@ -516,7 +527,7 @@ def test_message_file_to_stream_response_builds_signed_url(self, message_cycle_m ) cycle_db.commit() - with patch("core.app.task_pipeline.message_cycle_manager.sign_tool_file") as mock_sign: + with patch("services.app.generation.adapters.message_cycle.sign_tool_file") as mock_sign: mock_sign.return_value = "signed-url" response = message_cycle_manager.message_file_to_stream_response(SimpleNamespace(message_file_id="file-1")) @@ -550,44 +561,12 @@ def test_handle_retriever_resources_skips_none_entries(self, message_cycle_manag assert len(message_cycle_manager._task_state.metadata.retriever_resources) == 1 assert message_cycle_manager._task_state.metadata.retriever_resources[0].position == 1 - def test_knowledge_receipts_survive_message_metadata_without_collapsing_same_document(self, message_cycle_manager): - from dify_agent.protocol.knowledge_fs import KnowledgeFsCitation - - message_cycle_manager._application_generate_entity.app_config = _AppConfigWithFeatures( - additional_features=AppAdditionalFeatures(show_retrieve_source=True) - ) - message_cycle_manager._task_state = WorkflowTaskState(metadata=TaskStateMetadata()) - sources = [ - RetrievalSourceMetadata( - dataset_id="space", - document_id="doc", - data_source_type="knowledge_fs", - knowledge_fs_citation=KnowledgeFsCitation( - id="kfs_" + char * 32, - control_space_id="00000000-0000-4000-8000-000000000001", - space_name="Docs", - node_id=f"node-{char}", - document_asset_id="doc", - artifact_hash="hash", - ), - ) - for char in ("a", "b") - ] - event = QueueRetrieverResourcesEvent(retriever_resources=[*sources, sources[0]]) - message_cycle_manager.handle_retriever_resources(event) - message_cycle_manager.handle_retriever_resources(event) - restored = TaskStateMetadata.model_validate_json(message_cycle_manager._task_state.metadata.model_dump_json()) - assert [source.knowledge_fs_citation.id for source in restored.retriever_resources] == [ - "kfs_" + "a" * 32, - "kfs_" + "b" * 32, - ] - assert [source.position for source in restored.retriever_resources] == [1, 2] - def test_message_file_to_stream_response_uses_http_url_directly(self, message_cycle_manager, cycle_db: Session): """Use original URL when message file URL is already HTTP.""" message_cycle_manager._application_generate_entity.task_id = "task-http" cycle_db.add( _message_file( + cycle_db, file_id="file-http", message_id="msg-http", belongs_to=MessageFileBelongsTo.ASSISTANT, @@ -609,6 +588,7 @@ def test_message_file_to_stream_response_defaults_extension_to_bin_without_dot( message_cycle_manager._application_generate_entity.task_id = "task-bin" cycle_db.add( _message_file( + cycle_db, file_id="file-bin", message_id="msg-bin", belongs_to=MessageFileBelongsTo.ASSISTANT, @@ -618,7 +598,7 @@ def test_message_file_to_stream_response_defaults_extension_to_bin_without_dot( ) cycle_db.commit() - with patch("core.app.task_pipeline.message_cycle_manager.sign_tool_file") as mock_sign: + with patch("services.app.generation.adapters.message_cycle.sign_tool_file") as mock_sign: mock_sign.return_value = "signed-bin-url" response = message_cycle_manager.message_file_to_stream_response( @@ -645,3 +625,92 @@ def test_message_replace_to_stream_response_returns_reason(self, message_cycle_m assert response.answer == "replaced" assert response.reason == "moderation" + + def test_knowledge_receipts_survive_message_metadata_without_collapsing_same_document(self, message_cycle_manager): + from dify_agent.protocol.knowledge_fs import KnowledgeFsCitation + + message_cycle_manager._application_generate_entity.app_config = _AppConfigWithFeatures( + additional_features=AppAdditionalFeatures(show_retrieve_source=True) + ) + message_cycle_manager._task_state = WorkflowTaskState(metadata=TaskStateMetadata()) + sources = [ + RetrievalSourceMetadata( + dataset_id="space", + document_id="doc", + data_source_type="knowledge_fs", + knowledge_fs_citation=KnowledgeFsCitation( + id="kfs_" + char * 32, + control_space_id="00000000-0000-4000-8000-000000000001", + space_name="Docs", + node_id=f"node-{char}", + document_asset_id="doc", + artifact_hash="hash", + ), + ) + for char in ("a", "b") + ] + event = QueueRetrieverResourcesEvent(retriever_resources=[*sources, sources[0]]) + message_cycle_manager.handle_retriever_resources(event) + message_cycle_manager.handle_retriever_resources(event) + restored = TaskStateMetadata.model_validate_json(message_cycle_manager._task_state.metadata.model_dump_json()) + assert [source.knowledge_fs_citation.id for source in restored.retriever_resources] == [ + "kfs_" + "a" * 32, + "kfs_" + "b" * 32, + ] + assert [source.position for source in restored.retriever_resources] == [1, 2] + + +def test_conversation_name_model_call_holds_no_database_connection( + app, + app_records, + sqlite_session, + sqlite_engine, + monkeypatch, +): + from core.app.app_config.entities import WorkflowUIBasedAppConfig + from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom + from models.model import AppMode + from services.app.generation.adapters import message_cycle + + sqlite_session.add_all([_app(), _conversation()]) + sqlite_session.commit() + entity = AdvancedChatAppGenerateEntity( + task_id="task", + app_config=WorkflowUIBasedAppConfig( + tenant_id="tenant-1", + app_id="app-id", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow", + ), + inputs={}, + query="Hello", + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + ) + manager = MessageCycleManager( + records=app_records, application_generate_entity=entity, task_state=WorkflowTaskState() + ) + checked_out = set() + + def checkout(_dbapi, record, _proxy): + checked_out.add(id(record)) + + def checkin(_dbapi, record): + checked_out.discard(id(record)) + + def generate_name(*_args, **_kwargs): + assert checked_out == set() + return "Generated name" + + monkeypatch.setattr(message_cycle.LLMGenerator, "generate_conversation_name", generate_name) + event.listen(sqlite_engine, "checkout", checkout) + event.listen(sqlite_engine, "checkin", checkin) + try: + manager._generate_conversation_name_worker(app, "conv-1", "Hello") + assert checked_out == set() + finally: + event.remove(sqlite_engine, "checkout", checkout) + event.remove(sqlite_engine, "checkin", checkin) + assert sqlite_session.get(Conversation, "conv-1").name == "Generated name" diff --git a/api/tests/unit_tests/core/app/workflow/test_file_runtime.py b/api/tests/unit_tests/core/app/workflow/test_file_runtime.py index 294b3906f62f91..ac2209d77cd2a1 100644 --- a/api/tests/unit_tests/core/app/workflow/test_file_runtime.py +++ b/api/tests/unit_tests/core/app/workflow/test_file_runtime.py @@ -211,7 +211,7 @@ def test_resolve_upload_file_url_signs_internal_urls_and_supports_attachments( INTERNAL_FILES_URL="https://internal.example.com", ) monkeypatch.setattr("core.app.workflow.file_runtime.time.time", lambda: 1700000000) - monkeypatch.setattr("core.app.workflow.file_runtime.os.urandom", lambda _: b"\x01" * 16) + monkeypatch.setattr("libs.signed_query.os.urandom", lambda _: b"\x01" * 16) runtime = _build_runtime() url = runtime.resolve_upload_file_url( diff --git a/api/tests/unit_tests/core/workflow/nodes/human_input_v2/test_node.py b/api/tests/unit_tests/core/workflow/nodes/human_input_v2/test_node.py index 566eb71356e0f2..0ba57ededc5872 100644 --- a/api/tests/unit_tests/core/workflow/nodes/human_input_v2/test_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/human_input_v2/test_node.py @@ -55,6 +55,8 @@ StringSource, ) from repositories.human_input_v2.form_repository import Form, FormSubmission +from services.tools.provider_queries import ToolProviderIcons +from services.workflow.execution.ports import ExecutionContexts from tests.unit_tests.core.app.apps.common.test_workflow_response_converter import ( _workflow_response_converter_human_input_build_converter as _build_converter, ) @@ -269,7 +271,13 @@ def test_reentry_consumes_the_new_persisted_outcome(form: Form) -> None: @pytest.mark.parametrize("status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) -def test_actual_graph_queue_response_path_preserves_order_and_fields(form: Form, status: HumanInputFormStatus) -> None: +def test_actual_graph_queue_response_path_preserves_order_and_fields( + form: Form, + status: HumanInputFormStatus, + *, + workflow_contexts: ExecutionContexts, + tool_providers: ToolProviderIcons, +) -> None: form = submit(form) if status == HumanInputFormStatus.SUBMITTED else replace(form, status=status) node, _ = build_node(form) runner = create_autospec(WorkflowBasedAppRunner, instance=True) @@ -278,7 +286,7 @@ def test_actual_graph_queue_response_path_preserves_order_and_fields(form: Form, WorkflowBasedAppRunner.handle_event(runner, entry, event) queue_events = [call.args[0] for call in runner._publish_event.call_args_list] assert isinstance(queue_events[1], QueueNodeSucceededEvent) - converter = _build_converter() + converter = _build_converter(workflow_contexts=workflow_contexts, tool_providers=tool_providers) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL ) @@ -335,7 +343,9 @@ def test_dify_projection_serializes_the_runtime_token_without_a_graph_reason(for } -def test_file_segments_survive_node_graph_queue_and_response(form: Form) -> None: +def test_file_segments_survive_node_graph_queue_and_response( + form: Form, *, workflow_contexts: ExecutionContexts, tool_providers: ToolProviderIcons +) -> None: attachment = File( file_id="file-1", file_type=FileType.DOCUMENT, @@ -376,7 +386,7 @@ def test_file_segments_survive_node_graph_queue_and_response(form: Form) -> None queue_event = runner._publish_event.call_args.args[0] assert isinstance(queue_event.submitted_data["attachment"], FileSegment) assert isinstance(queue_event.submitted_data["files"], ArrayFileSegment) - converter = _build_converter() + converter = _build_converter(workflow_contexts=workflow_contexts, tool_providers=tool_providers) converter.workflow_start_to_stream_response( task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL ) diff --git a/api/tests/unit_tests/core/workflow/test_node_factory.py b/api/tests/unit_tests/core/workflow/test_node_factory.py index 124f25505d6078..9caeb262da4f42 100644 --- a/api/tests/unit_tests/core/workflow/test_node_factory.py +++ b/api/tests/unit_tests/core/workflow/test_node_factory.py @@ -20,7 +20,7 @@ from core.workflow import node_factory from core.workflow import template_rendering as workflow_template_rendering from core.workflow.llm_node import DifyLLMNode -from core.workflow.node_factory import DifyNodeFactory, get_node_type_classes_mapping +from core.workflow.node_factory import get_node_type_classes_mapping from core.workflow.node_runtime import DifyPreparedLLM from core.workflow.nodes.knowledge_index import KNOWLEDGE_INDEX_NODE_TYPE from graphon.entities.base_node_data import BaseNodeData @@ -42,26 +42,41 @@ from graphon.nodes.variable_assigner.v2.node import ( VariableAssignerNode as VariableAssignerV2, ) -from graphon.runtime import RuntimeState, VariablePool +from graphon.runtime import InitParams, RuntimeState, VariablePool from graphon.variables.segments import ArrayObjectSegment, ObjectSegment, StringSegment from models.base import TypeBase from models.model import AppMode, Conversation, ConversationFromSource +from services.workflow.execution.adapters.node_factory import DifyNodeFactory as ServiceDifyNodeFactory +from services.workflow.execution.ports import WorkflowRuntime from tests.helpers.unit.core.model_fixtures import make_model_instance from tests.workflow_test_utils import build_test_graph_init_params class DummyNode: - def __init__(self, *, node_id, data, init_params, runtime_state, **kwargs): + def __init__( + self, + *, + node_id: str, + data: dict[str, object], + init_params: InitParams, + runtime_state: RuntimeState, + **kwargs: object, + ) -> None: self.id = node_id self.data = data self.init_params = init_params self.runtime_state = runtime_state self.kwargs = kwargs + @classmethod + def validate_node_data(cls, node_data: BaseNodeData | Mapping[str, object]) -> BaseNodeData: + payload = node_data.model_dump(mode="python") if isinstance(node_data, BaseNodeData) else dict(node_data) + return BaseNodeData.model_validate(payload) + class DummyCodeNode(DummyNode): @classmethod - def default_code_providers(cls): + def default_code_providers(cls) -> tuple[()]: return () @@ -100,13 +115,13 @@ def _node_config(self, config_overrides: Callable[..., None]) -> None: ) @staticmethod - def _stub_node_resolution(monkeypatch: pytest.MonkeyPatch, node_class): + def _stub_node_resolution(monkeypatch: pytest.MonkeyPatch, node_class: type[DummyNode]) -> None: monkeypatch.setattr( - "core.workflow.node_factory.resolve_workflow_node_class", + "services.workflow.execution.adapters.node_factory.resolve_workflow_node_class", lambda **_kwargs: node_class, ) - def _factory(self): + def _factory(self) -> ServiceDifyNodeFactory: run_context = build_dify_run_context( tenant_id="tenant", app_id="app", @@ -115,47 +130,60 @@ def _factory(self): invoke_from=InvokeFrom.WEB_APP, ) - return DifyNodeFactory( - init_params=SimpleNamespace(run_context=run_context), - runtime_state=SimpleNamespace(), + return ServiceDifyNodeFactory( + init_params=InitParams( + workflow_id="workflow-id", + graph_config={}, + run_context=run_context, + call_depth=0, + ), + runtime_state=RuntimeState( + workflow_id="workflow-id", + variable_pool=VariablePool(), + start_at=0, + ), ) - def test_create_node_unknown_type(self): + def test_create_node_unknown_type(self) -> None: factory = self._factory() with pytest.raises(ValueError, match="No class mapping found for node type: unknown"): factory.create_node({"id": "node-1", "data": {"type": "unknown"}}) - def test_create_node_missing_mapping(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_missing_mapping(self, monkeypatch: pytest.MonkeyPatch) -> None: factory = self._factory() - monkeypatch.setattr("core.workflow.node_factory.get_node_type_classes_mapping", lambda: {}) + monkeypatch.setattr( + "services.workflow.execution.adapters.node_factory.get_node_type_classes_mapping", lambda: {} + ) with pytest.raises(ValueError, match="No class mapping found for node type: start"): factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START}}) - def test_create_node_missing_latest_class(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_missing_latest_class(self, monkeypatch: pytest.MonkeyPatch) -> None: factory = self._factory() monkeypatch.setattr( - "core.workflow.node_factory.get_node_type_classes_mapping", + "services.workflow.execution.adapters.node_factory.get_node_type_classes_mapping", lambda: {BuiltinNodeTypes.START: {"1": None}}, ) - monkeypatch.setattr("core.workflow.node_factory.LATEST_VERSION", "latest") + monkeypatch.setattr("services.workflow.execution.adapters.node_factory.LATEST_VERSION", "latest") with pytest.raises(ValueError, match="No latest version class found for node type: start"): factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START}}) - def test_create_node_selects_versioned_class(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_selects_versioned_class(self, monkeypatch: pytest.MonkeyPatch) -> None: factory = self._factory() selected_versions: list[tuple[str, str]] = [] class DummyNodeV2(DummyNode): pass - def _get_mapping(): + def _get_mapping() -> Mapping[NodeType, Mapping[str, type[DummyNode]]]: selected_versions.append(("snapshot", "called")) return {BuiltinNodeTypes.START: {"1": DummyNode, "2": DummyNodeV2}} - monkeypatch.setattr("core.workflow.node_factory.get_node_type_classes_mapping", _get_mapping) + monkeypatch.setattr( + "services.workflow.execution.adapters.node_factory.get_node_type_classes_mapping", _get_mapping + ) node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START, "version": "2"}}) @@ -163,7 +191,7 @@ def _get_mapping(): assert node.id == "node-1" assert selected_versions == [("snapshot", "called")] - def test_create_node_code_branch(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_code_branch(self, monkeypatch: pytest.MonkeyPatch) -> None: factory = self._factory() self._stub_node_resolution(monkeypatch, DummyCodeNode) @@ -172,7 +200,7 @@ def test_create_node_code_branch(self, monkeypatch: pytest.MonkeyPatch): assert isinstance(node, DummyCodeNode) assert node.id == "node-1" - def test_create_node_template_transform_branch(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_template_transform_branch(self, monkeypatch: pytest.MonkeyPatch) -> None: factory = self._factory() self._stub_node_resolution(monkeypatch, DummyTemplateTransformNode) @@ -181,7 +209,7 @@ def test_create_node_template_transform_branch(self, monkeypatch: pytest.MonkeyP assert isinstance(node, DummyTemplateTransformNode) assert "jinja2_template_renderer" in node.kwargs - def test_create_node_http_request_branch(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_http_request_branch(self, monkeypatch: pytest.MonkeyPatch) -> None: factory = self._factory() self._stub_node_resolution(monkeypatch, DummyHttpRequestNode) @@ -190,16 +218,19 @@ def test_create_node_http_request_branch(self, monkeypatch: pytest.MonkeyPatch): assert isinstance(node, DummyHttpRequestNode) assert "http_request_config" in node.kwargs - def test_create_node_knowledge_retrieval_branch(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_knowledge_retrieval_branch( + self, monkeypatch: pytest.MonkeyPatch, workflow_runtime: WorkflowRuntime + ) -> None: factory = self._factory() + factory._workflow_runtime = workflow_runtime self._stub_node_resolution(monkeypatch, DummyKnowledgeRetrievalNode) node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL}}) assert isinstance(node, DummyKnowledgeRetrievalNode) - assert node.kwargs == {} + assert "retrieval" in node.kwargs - def test_create_node_document_extractor_branch(self, monkeypatch: pytest.MonkeyPatch): + def test_create_node_document_extractor_branch(self, monkeypatch: pytest.MonkeyPatch) -> None: factory = self._factory() self._stub_node_resolution(monkeypatch, DummyDocumentExtractorNode) diff --git a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py index 54e54fe0dea60f..80e94b87f8a612 100644 --- a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py +++ b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py @@ -12,6 +12,7 @@ from opentelemetry.trace import StatusCode from sqlalchemy import Engine +from extensions.application_services.workflow import WorkflowExecutionDependencies from extensions.otel.decorators.handlers.workflow_app_runner_handler import WorkflowAppRunnerHandler from extensions.otel.semconv import DifySpanAttributes, GenAIAttributes from services.workflow_run_agg import WorkflowRunAgg @@ -19,8 +20,10 @@ from tests.unit_tests.services.test_workflow_run_agg import make_workflow_runner -def test_service_driver_emits_workflow_span(sqlite_engine: Engine) -> None: - runner, runs, nodes = make_workflow_runner(sqlite_engine) +def test_service_driver_emits_workflow_span( + sqlite_engine: Engine, workflow_runtime: WorkflowExecutionDependencies +) -> None: + runner, runs, nodes = make_workflow_runner(sqlite_engine, workflow_runtime) exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) diff --git a/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py b/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py index bdc93df83f5f0f..13c51405501a59 100644 --- a/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py +++ b/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py @@ -79,6 +79,8 @@ FormExecutionContext, HumanInputFormService, ) +from services.tools.provider_queries import ToolProviderIcons +from services.workflow.execution.ports import ExecutionContexts _TENANT = TenantId("00000000-0000-0000-0000-000000000001") _APP = "00000000-0000-0000-0000-000000000002" @@ -1517,7 +1519,12 @@ def send(candidate: ResendCandidate, *, to: str, subject: str, html: str) -> Non @pytest.mark.usefixtures("providers") @pytest.mark.parametrize("recipient", [Initiator(), OnetimeEmail(email="person@example.com")]) def test_pause_persistence_and_reconnect_preserve_v2_form_and_initiator_only_token( - context: Context, monkeypatch: pytest.MonkeyPatch, recipient: Initiator | OnetimeEmail + context: Context, + monkeypatch: pytest.MonkeyPatch, + recipient: Initiator | OnetimeEmail, + *, + workflow_contexts: ExecutionContexts, + tool_providers: ToolProviderIcons, ) -> None: from core.app.entities.queue_entities import QueueWorkflowPausedEvent from core.app.entities.task_entities import ( @@ -1600,7 +1607,7 @@ def reject_live_reload() -> Session: raise AssertionError("Unexpected live form reload") monkeypatch.setattr(session_factory, "create_session", reject_live_reload) - converter = _build_converter() + converter = _build_converter(workflow_contexts=workflow_contexts, tool_providers=tool_providers) converter.workflow_start_to_stream_response( task_id="task", workflow_run_id=_RUN, workflow_id=run.workflow_id, reason=WorkflowStartReason.INITIAL ) diff --git a/api/tests/unit_tests/services/test_workflow_node_variables.py b/api/tests/unit_tests/services/test_workflow_node_variables.py index 71a9e2b1295ce3..c6959ca52d449a 100644 --- a/api/tests/unit_tests/services/test_workflow_node_variables.py +++ b/api/tests/unit_tests/services/test_workflow_node_variables.py @@ -5,7 +5,6 @@ from sqlalchemy import Engine from sqlalchemy.orm import Session -from core.workflow.workflow_entry import WorkflowEntry from graphon.entities.graph_config import NodeConfigDictAdapter from graphon.nodes.list_operator.entities import FilterOperator from graphon.runtime import VariablePool @@ -13,6 +12,7 @@ from graphon.variables import StringSegment, VariableBase from models.human_input_entities import HumanInputNodeData from models.workflow import WorkflowDraftVariable +from services.workflow.execution.adapters.workflow_entry import WorkflowEntry from services.workflow_draft_variable_service import DraftVarLoader from services.workflow_node_variables import get_human_input_form_variable_mapping, load_additional_node_variables diff --git a/api/tests/unit_tests/services/test_workflow_persistence.py b/api/tests/unit_tests/services/test_workflow_persistence.py index cf508c3f14e86a..854d55c7d81ecf 100644 --- a/api/tests/unit_tests/services/test_workflow_persistence.py +++ b/api/tests/unit_tests/services/test_workflow_persistence.py @@ -70,15 +70,15 @@ from repositories.sqlalchemy_api_workflow_node_execution_repository import ( DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) -from services import workflow_persistence as persistence_mod from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer +from services.workflow.execution.adapters import persistence as persistence_mod +from services.workflow.execution.adapters.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from tests.helpers.unit.config_override import apply_config_overrides from tests.workflow_test_utils import build_test_graph_init_params """Verify the workflow persistence layer fans Inspector deltas to redis pub/sub. -The hook lives in ``services/workflow_persistence.py``: +The hook lives in ``services.workflow.execution.adapters.persistence``: every ``_handle_node_*`` and the terminal ``_handle_graph_run_*`` handlers call into ``services.workflow.inspector_events.publish_node_changed`` / ``publish_workflow_completed`` after the DB write succeeds. Those calls are @@ -921,7 +921,7 @@ def test_prepare_workflow_inputs_excludes_conversation_id(self, monkeypatch: pyt layer, _, _, _ = _make_layer() monkeypatch.setattr( - "core.workflow.workflow_entry.WorkflowEntry.handle_special_values", + "services.workflow.execution.adapters.workflow_entry.WorkflowEntry.handle_special_values", lambda inputs: inputs, ) diff --git a/api/tests/unit_tests/services/test_workflow_run_agg.py b/api/tests/unit_tests/services/test_workflow_run_agg.py index aa75358477bc54..68b3cb4b2f912b 100644 --- a/api/tests/unit_tests/services/test_workflow_run_agg.py +++ b/api/tests/unit_tests/services/test_workflow_run_agg.py @@ -1,17 +1,19 @@ import json +from dataclasses import replace from datetime import UTC, datetime, timedelta from pathlib import Path from uuid import uuid4 import pytest from sqlalchemy import Engine, select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from clients.agent_backend import FakeAgentBackendRunClient, FakeAgentBackendScenario +from core.app.apps.pipeline.pipeline_queue_manager import PipelineQueueManager from core.app.apps.pipeline.pipeline_runner import PipelineRunner -from core.app.apps.workflow.app_config_manager import WorkflowAppConfig -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow.app_runner import WorkflowAppRunner +from core.app.apps.workflow.app_config_manager import WorkflowAppConfig as CoreWorkflowAppConfig +from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager as CoreWorkflowAppQueueManager +from core.app.apps.workflow.app_runner import WorkflowAppRunner as CoreWorkflowAppRunner from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.app.entities.queue_entities import ( QueueNodeStartedEvent, @@ -24,6 +26,7 @@ from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from extensions.application_services.workflow import WorkflowExecutionDependencies from extensions.ext_storage import storage from extensions.storage.opendal_storage import OpenDALStorage from graphon.engine_events import ( @@ -39,7 +42,7 @@ from graphon.nodes.start.entities import StartNodeData from graphon.runtime import RuntimeState from graphon.variable_loader import DUMMY_VARIABLE_LOADER -from models import Account +from models.account import Account from models.agent import AgentWorkspace, AgentWorkspaceBinding, WorkflowAgentBindingType from models.dataset import Document from models.enums import AppTriggerType, CreatorUserRole, WorkflowRunTriggeredFrom, WorkflowTriggerStatus @@ -57,14 +60,18 @@ ) from models.workflow import WorkflowType as ModelWorkflowType from repositories.human_input_v2.form_repository import Form +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from repositories.workflow.definition_repository import workflow_snapshot from services.agent.workspace_service import AgentWorkspaceService from services.workflow.entities import WorkflowScheduleCFSPlanEntity +from services.workflow.execution.adapters.workflow.app_config_manager import WorkflowAppConfig +from services.workflow.execution.adapters.workflow.app_queue_manager import WorkflowAppQueueManager +from services.workflow.execution.adapters.workflow.app_runner import WorkflowAppRunner from services.workflow_run_agg import WorkflowRunAgg from tasks.workflow_cfs_scheduler.cfs_scheduler import AsyncWorkflowCFSPlanEntity from tasks.workflow_cfs_scheduler.entities import AsyncWorkflowQueue from tests.unit_tests.core.app.apps.pipeline.test_pipeline_runner import ( _build_app_generate_entity, - _build_runner, _document, _persist_scope, _workflow, @@ -77,7 +84,9 @@ def _pipeline_tool_runner( - sqlite_engine: Engine, invoke_from: InvokeFrom, source_node_data: dict[str, object] + sqlite_engine: Engine, + invoke_from: InvokeFrom, + source_node_data: dict[str, object], ) -> tuple[PipelineRunner, SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: source_graph = { "nodes": [ @@ -125,7 +134,7 @@ def _pipeline_tool_runner( ) session.add(provider) session.flush() - graph = { + graph: dict[str, object] = { "nodes": [ {"id": "start", "data": {"type": "start", "title": "Start", "variables": list[object]()}}, { @@ -149,7 +158,19 @@ def _pipeline_tool_runner( entity = _build_app_generate_entity().model_copy( update={"invoke_from": invoke_from, "document_id": None if invoke_from == InvokeFrom.DEBUGGER else "doc"} ) - runner = _build_runner(entity, sqlite_engine, workflow) + runner = PipelineRunner( + application_generate_entity=entity, + queue_manager=PipelineQueueManager( + task_id=entity.task_id, + user_id=entity.user_id, + invoke_from=entity.invoke_from, + app_mode=AppMode.RAG_PIPELINE, + ), + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=workflow, + system_user_id="sys", + documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), + ) user = Account(name="User", email="pipeline@example.com") user.id = entity.user_id runs = SQLAlchemyWorkflowExecutionRepository( @@ -284,7 +305,10 @@ def collect_task(**kwargs: object) -> None: def make_workflow_runner( - sqlite_engine: Engine, *, human_input: bool = False + sqlite_engine: Engine, + workflow_runtime: WorkflowExecutionDependencies, + *, + human_input: bool = False, ) -> tuple[WorkflowAppRunner, SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: """Build the same runner/repositories/queue used by a workflow worker.""" tenant_id, app_id, workflow_id, user_id, run_id = (str(uuid4()) for _ in range(5)) @@ -355,6 +379,90 @@ def make_workflow_runner( triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) runner = WorkflowAppRunner( + runtime=workflow_runtime, + application_generate_entity=entity, + queue_manager=queue, + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=workflow_snapshot(workflow), + system_user_id=user_id, + ) + return runner, workflow_repository, node_repository + + +def _make_legacy_workflow_runner( + sqlite_engine: Engine, *, human_input: bool = False +) -> tuple[ + CoreWorkflowAppRunner, + SQLAlchemyWorkflowExecutionRepository, + SQLAlchemyWorkflowNodeExecutionRepository, +]: + """Build the live core runner retained until atomic source activation migrates in #43680.""" + tenant_id, app_id, workflow_id, user_id, run_id = (str(uuid4()) for _ in range(5)) + graph: dict[str, list[dict[str, object]]] = { + "nodes": [{"id": "start", "data": {"type": "start", **StartNodeData(title="Start").model_dump(mode="json")}}], + "edges": [], + } + if human_input: + data = HumanInputNodeData( + title="Approval", + form_content="Decision: {{#$output.answer#}}", + inputs=[ParagraphInputConfig(output_variable_name="answer")], + user_actions=[UserActionConfig(id="approve", title="Approve")], + ) + graph["nodes"].append({"id": "approval", "data": {"type": "human-input", **data.model_dump(mode="json")}}) + graph["edges"].append({"id": "start-approval", "source": "start", "target": "approval"}) + graph["nodes"].append( + { + "id": "end", + "data": { + "type": "end", + "title": "End", + "outputs": [{"variable": "answer", "value_selector": ["approval", "answer"]}], + }, + } + ) + graph["edges"].append({"id": "approval-end", "source": "approval", "target": "end", "sourceHandle": "approve"}) + workflow = Workflow( + id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + type=ModelWorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph), + _features="{}", + created_by=user_id, + ) + entity = WorkflowAppGenerateEntity( + user_from=UserFrom.ACCOUNT, + task_id=str(uuid4()), + app_config=CoreWorkflowAppConfig( + tenant_id=tenant_id, app_id=app_id, workflow_id=workflow_id, app_mode=AppMode.WORKFLOW + ), + inputs={}, + files=[], + user_id=user_id, + stream=True, + invoke_from=InvokeFrom.DEBUGGER, + workflow_execution_id=run_id, + ) + user = Account(name="Workflow Tester", email=f"{user_id}@example.com") + user.id = user_id + queue = CoreWorkflowAppQueueManager(entity.task_id, user_id, entity.invoke_from, AppMode.WORKFLOW) + workflow_repository = SQLAlchemyWorkflowExecutionRepository( + sqlite_engine, + tenant_id=tenant_id, + user=user, + app_id=app_id, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + ) + node_repository = SQLAlchemyWorkflowNodeExecutionRepository( + sqlite_engine, + tenant_id=tenant_id, + user=user, + app_id=app_id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + runner = CoreWorkflowAppRunner( application_generate_entity=entity, queue_manager=queue, variable_loader=DUMMY_VARIABLE_LOADER, @@ -367,8 +475,10 @@ def make_workflow_runner( class TestWorkflowRunAggCases: @staticmethod - def test_driver_publishes_persisted_start_index_and_success(sqlite_engine: Engine) -> None: - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine) + def test_driver_publishes_persisted_start_index_and_success( + sqlite_engine: Engine, workflow_runtime: WorkflowExecutionDependencies + ) -> None: + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, workflow_runtime) WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) events = [message.event for message in runner._queue_manager.listen()] starts = [event for event in events if isinstance(event, QueueNodeStartedEvent)] @@ -384,10 +494,12 @@ def test_driver_publishes_persisted_start_index_and_success(sqlite_engine: Engin @staticmethod @pytest.mark.parametrize("snapshot", [None, {}, {"tools": []}]) - def test_invalid_saved_tool_snapshot_is_rejected(sqlite_engine: Engine, snapshot: object) -> None: + def test_invalid_saved_tool_snapshot_is_rejected( + sqlite_engine: Engine, workflow_runtime: WorkflowExecutionDependencies, snapshot: object + ) -> None: from pydantic import ValidationError - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine) + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, workflow_runtime) runner.application_generate_entity.extras["workflow_tool_snapshot"] = snapshot with pytest.raises(ValidationError): WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) @@ -398,19 +510,22 @@ def test_invalid_saved_tool_snapshot_is_rejected(sqlite_engine: Engine, snapshot @pytest.fixture def _local_storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(storage, "storage_runner", OpenDALStorage(scheme="fs", root=str(tmp_path)), raising=False) - from core.app.apps import workflow_app_runner + from services.workflow.execution.adapters import events - monkeypatch.setattr( - workflow_app_runner.dispatch_human_input_form_delivery_task, "apply_async", lambda **_kwargs: None - ) + monkeypatch.setattr(events.dispatch_human_input_form_delivery_task, "apply_async", lambda **_kwargs: None) @staticmethod @pytest.mark.usefixtures("_local_storage") @pytest.mark.parametrize("persist_snapshot", [False, True]) def test_pause_publication_persists_summary_and_optional_snapshot( - sqlite_engine: Engine, persist_snapshot: bool, caplog: pytest.LogCaptureFixture + sqlite_engine: Engine, + workflow_runtime: WorkflowExecutionDependencies, + persist_snapshot: bool, + caplog: pytest.LogCaptureFixture, ) -> None: - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + runner, workflow_repository, node_repository = make_workflow_runner( + sqlite_engine, workflow_runtime, human_input=True + ) config = ( PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id) if persist_snapshot @@ -442,7 +557,7 @@ def test_pause_publication_persists_summary_and_optional_snapshot( @staticmethod @pytest.mark.usefixtures("_local_storage") def test_pause_database_failure_publishes_failure_instead_of_resume_ready( - sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch + sqlite_engine: Engine, workflow_runtime: WorkflowExecutionDependencies, monkeypatch: pytest.MonkeyPatch ) -> None: def fail_pause(*_args: object, **_kwargs: object) -> None: raise RuntimeError("pause-write-failed") @@ -466,7 +581,9 @@ def retire(_store: object, **kwargs: object) -> list[str]: "core.workflow.nodes.agent_v2.workspace_retirement_layer.enqueue_agent_resource_collection", lambda **kwargs: collections.append(kwargs), ) - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + runner, workflow_repository, node_repository = make_workflow_runner( + sqlite_engine, workflow_runtime, human_input=True + ) trigger_log = WorkflowTriggerLog( tenant_id=runner.application_generate_entity.app_config.tenant_id, app_id=runner.application_generate_entity.app_config.app_id, @@ -532,13 +649,15 @@ def retire(_store: object, **kwargs: object) -> list[str]: assert session.scalar(select(WorkflowPause)) is None @staticmethod - def test_late_form_submission_is_refreshed_before_the_real_queue_closes(sqlite_session: Session) -> None: - from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner + def test_late_form_submission_is_refreshed_before_the_real_queue_closes( + sqlite_session: Session, workflow_runtime: WorkflowExecutionDependencies + ) -> None: from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent from core.repositories.human_input_repository import HumanInputFormSubmissionRepository from enums.human_input import HumanInputFormStatus from graphon.engine_events import NodeRunStartedEvent from tests.unit_tests.core.app.apps.test_workflow_app_runner import ( + _build_runner, _make_paused_workflow, _resume_events, _save_form, @@ -547,7 +666,7 @@ def test_late_form_submission_is_refreshed_before_the_real_queue_closes(sqlite_s first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) late = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) queue = WorkflowAppQueueManager("form-refresh-task", "user", InvokeFrom.DEBUGGER, AppMode.WORKFLOW) - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + runner = _build_runner(queue, workflow_runtime) entry = _make_paused_workflow(runner, [first, late]) for event in _resume_events(runner, entry): runner.handle_event(entry, event) @@ -568,12 +687,16 @@ def test_late_form_submission_is_refreshed_before_the_real_queue_closes(sqlite_s @staticmethod @pytest.mark.usefixtures("_local_storage") - def test_durable_pause_resumes_with_one_form_completion_and_continuing_indexes(sqlite_engine: Engine) -> None: + def test_durable_pause_resumes_with_one_form_completion_and_continuing_indexes( + sqlite_engine: Engine, workflow_runtime: WorkflowExecutionDependencies + ) -> None: from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent, QueueWorkflowStartedEvent from core.repositories.human_input_repository import HumanInputFormSubmissionRepository from models.human_input import HumanInputForm - original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + original, workflow_repository, node_repository = make_workflow_runner( + sqlite_engine, workflow_runtime, human_input=True + ) config = PauseStateLayerConfig(sqlite_engine, original.application_generate_entity.user_id) WorkflowRunAgg.run(original, config, workflow_repository, node_repository) before = [message.event for message in original._queue_manager.listen()] @@ -596,6 +719,7 @@ def test_durable_pause_resumes_with_one_form_completion_and_continuing_indexes(s entity = original.application_generate_entity.model_copy(update={"task_id": str(uuid4())}) queue = WorkflowAppQueueManager(entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW) resumed = WorkflowAppRunner( + runtime=workflow_runtime, application_generate_entity=entity, queue_manager=queue, variable_loader=DUMMY_VARIABLE_LOADER, @@ -603,7 +727,6 @@ def test_durable_pause_resumes_with_one_form_completion_and_continuing_indexes(s system_user_id=entity.user_id, graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), response_stream_filter=snapshot.get_response_stream_filter(), - execution_driver=WorkflowRunAgg.run, ) WorkflowRunAgg.run(resumed, config, workflow_repository, node_repository) after = [message.event for message in queue.listen()] @@ -631,7 +754,7 @@ def test_pause_before_first_tool_call_preserves_its_published_definition(sqlite_ from models.model import App from models.tools import WorkflowToolProvider - original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + original, workflow_repository, node_repository = _make_legacy_workflow_runner(sqlite_engine, human_input=True) original_entity = original.application_generate_entity tenant_id = original_entity.app_config.tenant_id source_app_id = str(uuid4()) @@ -757,8 +880,8 @@ def test_pause_before_first_tool_call_preserves_its_published_definition(sqlite_ ) entity = snapshot.get_generate_entity().model_copy(update={"task_id": str(uuid4())}) assert isinstance(entity, WorkflowAppGenerateEntity) - queue = WorkflowAppQueueManager(entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW) - resumed = WorkflowAppRunner( + queue = CoreWorkflowAppQueueManager(entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW) + resumed = CoreWorkflowAppRunner( application_generate_entity=entity, queue_manager=queue, variable_loader=DUMMY_VARIABLE_LOADER, @@ -788,10 +911,10 @@ def test_pause_before_first_tool_call_preserves_its_published_definition(sqlite_ fresh_entity = original_entity.model_copy( update={"task_id": str(uuid4()), "workflow_execution_id": str(uuid4()), "extras": {}} ) - fresh_queue = WorkflowAppQueueManager( + fresh_queue = CoreWorkflowAppQueueManager( fresh_entity.task_id, fresh_entity.user_id, fresh_entity.invoke_from, AppMode.WORKFLOW ) - fresh = WorkflowAppRunner( + fresh = CoreWorkflowAppRunner( application_generate_entity=fresh_entity, queue_manager=fresh_queue, variable_loader=DUMMY_VARIABLE_LOADER, @@ -805,10 +928,14 @@ def test_pause_before_first_tool_call_preserves_its_published_definition(sqlite_ assert fresh_events[-1].outputs == {"result": [{"marker": "republished"}]} @staticmethod - def test_active_execution_snapshot_failure_is_not_published_as_a_pause(sqlite_engine: Engine) -> None: + def test_active_execution_snapshot_failure_is_not_published_as_a_pause( + sqlite_engine: Engine, workflow_runtime: WorkflowExecutionDependencies + ) -> None: from graphon.engine_events import GraphRunFailedEvent, GraphRunPausedEvent - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + runner, workflow_repository, node_repository = make_workflow_runner( + sqlite_engine, workflow_runtime, human_input=True + ) prepared = runner.prepare() assert prepared is not None aggregate = WorkflowRunAgg( @@ -844,30 +971,33 @@ def test_active_execution_snapshot_failure_is_not_published_as_a_pause(sqlite_en def test_elapsed_time_accumulates_across_pauses_and_matches_delayed_stream_responses( monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, + workflow_runtime: WorkflowExecutionDependencies, terminal_event: GraphRunSucceededEvent | GraphRunPartialSucceededEvent | GraphRunFailedEvent | GraphRunAbortedEvent, status: WorkflowExecutionStatus, ) -> None: - from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter from core.app.entities.task_entities import WorkflowPauseStreamResponse from core.repositories.human_input_repository import HumanInputFormSubmissionRepository from core.workflow.system_variables import build_system_variables from enums.human_input import HumanInputFormStatus from graphon.entities import WorkflowStartReason from models.human_input import HumanInputForm + from services.workflow.execution.adapters.response_converter import WorkflowResponseConverter clock = [100.0] monkeypatch.setattr("repositories.entities.workflow_execution.perf_counter", lambda: clock[0]) - original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + original, workflow_repository, node_repository = make_workflow_runner( + sqlite_engine, workflow_runtime, human_input=True + ) graph = original._workflow.graph_dict graph["nodes"].append({"id": "second-approval", "data": graph["nodes"][1]["data"]}) graph["edges"][1]["target"] = "second-approval" graph["edges"].append( {"id": "second-end", "source": "second-approval", "target": "end", "sourceHandle": "approve"} ) - original._workflow.graph = json.dumps(graph) + original._workflow = replace(original._workflow, graph=json.dumps(graph)) runner = original run_id = runner.application_generate_entity.workflow_execution_id config = PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id) @@ -917,6 +1047,8 @@ def test_elapsed_time_accumulates_across_pauses_and_matches_delayed_stream_respo application_generate_entity=entity, user=account, system_variables=build_system_variables(workflow_execution_id=run_id), + contexts=workflow_runtime.contexts, + tool_providers=workflow_runtime.tool_providers, ) converter.workflow_start_to_stream_response( task_id=entity.task_id, @@ -956,6 +1088,7 @@ def test_elapsed_time_accumulates_across_pauses_and_matches_delayed_stream_respo ) entity.task_id = str(uuid4()) runner = WorkflowAppRunner( + runtime=workflow_runtime, application_generate_entity=entity, queue_manager=WorkflowAppQueueManager( entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW @@ -965,7 +1098,6 @@ def test_elapsed_time_accumulates_across_pauses_and_matches_delayed_stream_respo system_user_id=entity.user_id, graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), response_stream_filter=snapshot.get_response_stream_filter(), - execution_driver=WorkflowRunAgg.run, ) else: response = converter.workflow_finish_to_stream_response( @@ -981,7 +1113,10 @@ def test_elapsed_time_accumulates_across_pauses_and_matches_delayed_stream_respo @pytest.mark.usefixtures("_local_storage") @pytest.mark.parametrize("stop_timing", ["before_summary", "during_upload"]) def test_stop_wins_while_engine_is_persisting_a_pause( - sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch, stop_timing: str + sqlite_engine: Engine, + workflow_runtime: WorkflowExecutionDependencies, + monkeypatch: pytest.MonkeyPatch, + stop_timing: str, ) -> None: from sqlalchemy.orm import sessionmaker @@ -991,7 +1126,9 @@ def test_stop_wins_while_engine_is_persisting_a_pause( from models.human_input import HumanInputForm from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + runner, workflow_repository, node_repository = make_workflow_runner( + sqlite_engine, workflow_runtime, human_input=True + ) entity = runner.application_generate_entity repository = DifyAPISQLAlchemyWorkflowRunRepository(sessionmaker(sqlite_engine)) stop_calls = [] @@ -1046,7 +1183,7 @@ def persist(aggregate: WorkflowRunAgg, event: GraphRunPausedEvent) -> None: def test_legacy_caller_rejects_human_input_before_creating_run_or_form(sqlite_engine: Engine) -> None: from models.human_input import HumanInputForm - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + runner, workflow_repository, node_repository = _make_legacy_workflow_runner(sqlite_engine, human_input=True) runner.application_generate_entity.allow_human_input = False with pytest.raises(ValueError, match="Human Input nodes are not supported"): WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) @@ -1060,7 +1197,12 @@ def test_legacy_caller_rejects_human_input_before_creating_run_or_form(sqlite_en @pytest.mark.parametrize("token", ["initiator-token", None]) def test_v2_pause_persists_only_checkpoint_references_and_publishes_frozen_form( - sqlite_engine: Engine, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, token: str | None, v2_form: Form + sqlite_engine: Engine, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + token: str | None, + v2_form: Form, + workflow_runtime: WorkflowExecutionDependencies, ) -> None: from dataclasses import replace @@ -1072,7 +1214,11 @@ def test_v2_pause_persists_only_checkpoint_references_and_publishes_frozen_form( from models.workflow import WorkflowPauseReason monkeypatch.setattr(storage, "storage_runner", OpenDALStorage(scheme="fs", root=str(tmp_path)), raising=False) - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + runner, workflow_repository, node_repository = make_workflow_runner( + sqlite_engine, + workflow_runtime, + human_input=True, + ) entity = runner.application_generate_entity graph = runner._workflow.graph_dict graph["nodes"][1]["data"] = V2NodeData.model_validate( @@ -1087,7 +1233,7 @@ def test_v2_pause_persists_only_checkpoint_references_and_publishes_frozen_form( "debug_mode": {"enabled": False, "channels": []}, } ).model_dump(mode="json") - runner._workflow.graph = json.dumps(graph) + runner._workflow = replace(runner._workflow, graph=json.dumps(graph)) source_form = v2_form prepared_forms: list[PreparedForm] = [] @@ -1111,10 +1257,7 @@ def prepare(_runtime: object, *, node_execution_id: str, **_kwargs: object) -> P def reject_legacy_lookup(*_args: object, **_kwargs: object) -> None: pytest.fail("V2 forms must not be loaded from the legacy repository") - monkeypatch.setattr( - "core.repositories.human_input_repository.HumanInputFormSubmissionRepository.get_by_form_ids", - reject_legacy_lookup, - ) + monkeypatch.setattr(runner.human_form_reader, "get_by_form_id", reject_legacy_lookup) WorkflowRunAgg.run( runner, PauseStateLayerConfig(sqlite_engine, entity.user_id), workflow_repository, node_repository ) diff --git a/api/tests/unit_tests/services/test_workflow_run_index.py b/api/tests/unit_tests/services/test_workflow_run_index.py index dffb8e5fe6eeb6..121ee501e860ca 100644 --- a/api/tests/unit_tests/services/test_workflow_run_index.py +++ b/api/tests/unit_tests/services/test_workflow_run_index.py @@ -8,7 +8,6 @@ from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.node_factory import DifyNodeFactory from core.workflow.system_variables import build_system_variables from graphon.engine import Engine from graphon.engine_events import NodeRunStartedEvent @@ -16,9 +15,11 @@ from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowType from graphon.graph import Graph from graphon.runtime import RuntimeState, VariablePool -from models import Account, AppMode +from models.account import Account from models.enums import WorkflowRunTriggeredFrom +from models.model import AppMode from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom +from services.workflow.execution.adapters.node_factory import DifyNodeFactory from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from services.workflow_run_index import WorkflowRunIndex from tests.workflow_test_utils import build_test_graph_init_params diff --git a/api/tests/unit_tests/services/tracing/test_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index 775662e814e13b..b0a4ffc06d4ec2 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -21,6 +21,7 @@ from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom from models.model import App, AppMode, Message from models.workflow import WorkflowRun, WorkflowType +from repositories.app.generation_repository import AppGenerationRepository from services.tracing import recovery from services.tracing import service as service_module from services.tracing.normalization import span_id @@ -430,6 +431,7 @@ def test_stop_save_preserves_worker_usage_committed_after_its_cached_message_rea def test_stopped_runner_hands_off_only_an_observed_runtime_terminal( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, + app_records: AppGenerationRepository, runtime_outcome: str, final_batch_arrived: bool, snapshot_save_fails: bool, @@ -481,7 +483,7 @@ def unexpected_cancel(*_args: object, **_kwargs: object) -> None: monkeypatch.setattr(client, "cancel_run_and_wait", unexpected_cancel) store = _ExplodingSessionStore() if snapshot_save_fails else _FakeSessionStore() - runner = _runner(client, store) + runner = _runner(client, store, app_records=app_records) run = case.collection.run context = DifyRunContext( tenant_id=run.tenant_id, @@ -584,7 +586,10 @@ def test_missing_runtime_stop_terminal_converges_without_inventing_success( @pytest.mark.parametrize("stop_on_end", [False, True]) def test_stop_during_final_queue_publication_preserves_the_already_observed_success( - case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, stop_on_end: bool + case: RecoveryCase, + monkeypatch: pytest.MonkeyPatch, + app_records: AppGenerationRepository, + stop_on_end: bool, ) -> None: from cachetools import TTLCache from flask import Flask @@ -615,7 +620,7 @@ def test_stop_during_final_queue_publication_preserves_the_already_observed_succ base_app_queue_manager.redis_client, "setex", lambda key, _ttl, value: flags.__setitem__(key, value) ) monkeypatch.setattr( - "core.app.apps.agent_app.runtime_request_builder.resolve_model_context_window", lambda **_kwargs: None + "services.app.generation.adapters.agent_request_builder.resolve_model_context_window", lambda **_kwargs: None ) task_id = identifier("stop-at-final-publish") queue_manager = MessageBasedAppQueueManager( @@ -629,7 +634,7 @@ def test_stop_during_final_queue_publication_preserves_the_already_observed_succ clock = [0.0] queue_manager._stopped_cache = TTLCache(maxsize=1, ttl=1, timer=lambda: clock[0]) store = _FakeSessionStore() - runner = _runner(FakeAgentBackendRunClient(), store) + runner = _runner(FakeAgentBackendRunClient(), store, app_records=app_records) original_publish = queue_manager._publish def publish_after_stop(event: AppQueueEvent, pub_from: PublishFrom) -> None: @@ -681,7 +686,10 @@ def publish_after_stop(event: AppQueueEvent, pub_from: PublishFrom) -> None: @pytest.mark.parametrize("stop_during_cleanup", [False, True]) def test_runtime_cleanup_cancellation_does_not_reclassify_a_business_stream_failure( - case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, stop_during_cleanup: bool + case: RecoveryCase, + monkeypatch: pytest.MonkeyPatch, + app_records: AppGenerationRepository, + stop_during_cleanup: bool, ) -> None: from collections.abc import Iterator from typing import cast @@ -722,7 +730,7 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: monkeypatch.setattr(client, "cancel_run_and_wait", cancel_then_stop) store = _FakeSessionStore() - runner = _runner(client, store) + runner = _runner(client, store, app_records=app_records) run = case.collection.run context = DifyRunContext( tenant_id=run.tenant_id,