opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-24 23:22:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(workflow): support human input in loop and iteration (#39243)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
-LAN-
2026-08-03 02:24:18 +00:00
committed by GitHub
co-authored by autofix-ci[bot]
parent 8dbac96621
commit 351577bdb0
86 changed files with 1454 additions and 762 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/tests/integration_tests/workflow/nodes/test_tool.py
+1
View File
@@ -70,6 +70,7 @@ def init_tool_node(config: dict):
tool_file_manager=tool_file_manager,
runtime=DifyToolNodeRuntime(init_params.run_context),
)
node.bind_execution_id(str(uuid.uuid4()))
return node
api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py
+5 -11
View File
@@ -17,7 +17,6 @@ These tests use TestContainers to spin up real services for integration testing,
providing more reliable and realistic test scenarios than mocks.
"""
import json
import uuid
from time import time
from unittest.mock import Mock
@@ -237,12 +236,12 @@ class TestPauseStatePersistenceLayerTestContainers:
# Create LLM usage
llm_usage = LLMUsage.empty_usage()
llm_usage.total_tokens = total_tokens
# Create graph runtime state
graph_runtime_state = GraphRuntimeState(
variable_pool=variable_pool,
start_at=start_at,
total_tokens=total_tokens,
llm_usage=llm_usage,
outputs=outputs or {},
node_run_steps=node_run_steps,
@@ -366,9 +365,6 @@ class TestPauseStatePersistenceLayerTestContainers:
resumption_context = WorkflowResumptionContext.loads(storage_content)
assert resumption_context.version == "1"
assert resumption_context.serialized_graph_runtime_state == graph_runtime_state.dumps()
expected_state = json.loads(graph_runtime_state.dumps())
actual_state = json.loads(resumption_context.serialized_graph_runtime_state)
assert actual_state == expected_state
persisted_entity = resumption_context.get_generate_entity()
assert isinstance(persisted_entity, WorkflowAppGenerateEntity)
assert persisted_entity.workflow_execution_id == self.test_workflow_run_id
@@ -414,13 +410,11 @@ class TestPauseStatePersistenceLayerTestContainers:
state_bytes = pause_entity.get_state()
resumption_context = WorkflowResumptionContext.loads(state_bytes.decode())
retrieved_state = json.loads(resumption_context.serialized_graph_runtime_state)
expected_state = json.loads(graph_runtime_state.dumps())
retrieved_state = GraphRuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state)
assert retrieved_state == expected_state
assert retrieved_state["outputs"] == complex_outputs
assert retrieved_state["total_tokens"] == 250
assert retrieved_state["node_run_steps"] == 10
assert retrieved_state.outputs == complex_outputs
assert retrieved_state.total_tokens == 250
assert retrieved_state.node_run_steps == 10
assert resumption_context.get_generate_entity().workflow_execution_id == self.test_workflow_run_id
def test_database_transaction_handling(self, db_session_with_containers: Session):
api/tests/test_containers_integration_tests/models/test_types_enum_text.py
+6 -8
View File
@@ -210,7 +210,7 @@ class TestEnumText:
assert str(exc.value) == "'invalid' is not a valid _UserType"
def test_select_legacy_model_type_values(self, engine_with_containers: Engine):
def test_select_rejects_legacy_model_type_values(self, engine_with_containers: Engine):
insertion_sql = """
INSERT INTO enum_text_legacy_model_type_test (id, model_type) VALUES
(1, 'text-generation'),
@@ -221,11 +221,9 @@ class TestEnumText:
session.execute(sa.text(insertion_sql))
session.commit()
with Session(engine_with_containers) as session:
records = session.scalars(select(_LegacyModelTypeRecord).order_by(_LegacyModelTypeRecord.id)).all()
for record_id, legacy_value in enumerate(("text-generation", "embeddings", "reranking"), 1):
with pytest.raises(ValueError) as exc:
with Session(engine_with_containers) as session:
session.scalar(select(_LegacyModelTypeRecord).where(_LegacyModelTypeRecord.id == record_id))
assert [record.model_type for record in records] == [
ModelType.LLM,
ModelType.TEXT_EMBEDDING,
ModelType.RERANK,
]
assert str(exc.value) == f"'{legacy_value}' is not a valid ModelType"
api/tests/unit_tests/configs/test_dify_config.py
+1
View File
@@ -83,6 +83,7 @@ def test_dify_config(monkeypatch: pytest.MonkeyPatch):
assert config.AGENT_SHELL_ENABLED is True
assert config.SENTRY_TRACES_SAMPLE_RATE == 1.0
assert config.TEMPLATE_TRANSFORM_MAX_LENGTH == 400_000
assert config.GRAPH_ENGINE_SCALE_UP_THRESHOLD == 0
# annotated field with custom configured value
assert config.HTTP_REQUEST_MAX_READ_TIMEOUT == 300
api/tests/unit_tests/controllers/console/test_human_input_form.py
+61 -1
View File
@@ -4,7 +4,7 @@ import json
from datetime import UTC, datetime
from inspect import unwrap
from types import SimpleNamespace
from unittest.mock import Mock
from unittest.mock import ANY, Mock
import pytest
from flask import Flask, Response
@@ -18,6 +18,7 @@ from controllers.console.human_input_form import (
WorkflowResponseConverter,
_jsonify_form_definition,
)
from core.workflow.human_input_policy import HumanInputSurface
from models.account import AccountStatus
from models.enums import CreatorUserRole
from models.human_input import RecipientType
@@ -344,3 +345,62 @@ def test_workflow_events_finished(app: Flask, monkeypatch: pytest.MonkeyPatch) -
assert response.mimetype == "text/event-stream"
assert "data" in response.get_data(as_text=True)
def test_workflow_events_snapshot_can_continue_across_pauses(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
workflow_run = SimpleNamespace(
id="run-1",
created_by_role=CreatorUserRole.ACCOUNT,
created_by="user-1",
tenant_id="t1",
app_id="app-1",
finished_at=None,
)
app_model = SimpleNamespace(mode=AppMode.WORKFLOW)
class _RepoStub:
def get_workflow_run_by_id_and_tenant_id(self, **_kwargs):
return workflow_run
workflow_generator = Mock()
workflow_generator.convert_to_event_stream.return_value = iter(["data: snapshot\n\n"])
snapshot_builder = Mock(return_value=["snapshot-events"])
monkeypatch.setattr(
DifyAPIRepositoryFactory,
"create_api_workflow_run_repository",
lambda *_args, **_kwargs: _RepoStub(),
)
monkeypatch.setattr(
"controllers.console.human_input_form._retrieve_app_for_workflow_run",
lambda *_args, **_kwargs: app_model,
)
monkeypatch.setattr(
"controllers.console.human_input_form.WorkflowAppGenerator",
lambda: workflow_generator,
)
monkeypatch.setattr(
"controllers.console.human_input_form.build_workflow_event_stream",
snapshot_builder,
)
monkeypatch.setattr("controllers.console.human_input_form.db", SimpleNamespace(engine=object()))
api = ConsoleWorkflowEventsApi()
handler = unwrap(api.get)
with app.test_request_context(
"/console/api/workflow/run-1/events?include_state_snapshot=true&continue_on_pause=true",
method="GET",
):
response = handler(api, "t1", SimpleNamespace(id="user-1"), workflow_run_id="run-1")
assert response.get_data(as_text=True) == "data: snapshot\n\n"
snapshot_builder.assert_called_once_with(
app_mode=AppMode.WORKFLOW,
workflow_run=workflow_run,
tenant_id="t1",
app_id="app-1",
session_maker=ANY,
human_input_surface=HumanInputSurface.CONSOLE,
close_on_pause=False,
)
api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py
+1 -2
View File
@@ -271,8 +271,7 @@ def _build_resumption_context(task_id: str) -> WorkflowResumptionContext:
workflow_execution_id="run-1",
)
runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0)
runtime_state.register_paused_node("node-1")
runtime_state.outputs = {"result": "value"}
runtime_state.set_output("result", "value")
wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity)
return WorkflowResumptionContext(
generate_entity=wrapper,
api/tests/unit_tests/controllers/web/test_workflow_events.py
+38 -1
View File
@@ -3,7 +3,7 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from unittest.mock import ANY, MagicMock, Mock, patch
import pytest
from flask import Flask
@@ -11,6 +11,7 @@ from flask import Flask
from controllers.common.errors import NotFoundError
from controllers.web.workflow_events import WorkflowEventsApi
from models.enums import CreatorUserRole
from models.model import AppMode
def _workflow_app() -> SimpleNamespace:
@@ -125,3 +126,39 @@ class TestWorkflowEventsApi:
response = WorkflowEventsApi().get(_workflow_app(), _end_user(), "run-1")
assert response.mimetype == "text/event-stream"
@patch("controllers.web.workflow_events.DifyAPIRepositoryFactory")
@patch("controllers.web.workflow_events.db")
def test_snapshot_stream_can_continue_across_pauses(
self, mock_db: MagicMock, mock_factory: MagicMock, app: Flask, monkeypatch: pytest.MonkeyPatch
) -> None:
mock_db.engine = "engine"
run = SimpleNamespace(
id="run-1",
app_id="app-1",
created_by_role=CreatorUserRole.END_USER,
created_by="eu-1",
finished_at=None,
)
mock_repo = MagicMock()
mock_repo.get_workflow_run_by_id_and_tenant_id.return_value = run
mock_factory.create_api_workflow_run_repository.return_value = mock_repo
workflow_generator = Mock()
workflow_generator.convert_to_event_stream.return_value = iter(["data: snapshot\n\n"])
snapshot_builder = Mock(return_value=["snapshot-events"])
monkeypatch.setattr("controllers.web.workflow_events.WorkflowAppGenerator", lambda: workflow_generator)
monkeypatch.setattr("controllers.web.workflow_events.build_workflow_event_stream", snapshot_builder)
with app.test_request_context("/workflow/run-1/events?include_state_snapshot=true&continue_on_pause=true"):
response = WorkflowEventsApi().get(_workflow_app(), _end_user(), "run-1")
assert response.get_data(as_text=True) == "data: snapshot\n\n"
snapshot_builder.assert_called_once_with(
app_mode=AppMode.WORKFLOW,
workflow_run=run,
tenant_id="tenant-1",
app_id="app-1",
session_maker=ANY,
close_on_pause=False,
)
api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py
+8 -2
View File
@@ -748,11 +748,13 @@ class TestAdvancedChatAppGeneratorInternals:
monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context)
workflow = SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app")
class _Session:
def __init__(self, *args, **kwargs):
self.scalar = MagicMock(
side_effect=[
SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app"),
workflow,
SimpleNamespace(id="app"),
]
)
@@ -772,6 +774,8 @@ class TestAdvancedChatAppGeneratorInternals:
monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session)
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=object(), session=SimpleNamespace(close=lambda: None)),
@@ -788,10 +792,12 @@ class TestAdvancedChatAppGeneratorInternals:
workflow_execution_repository=SimpleNamespace(),
workflow_node_execution_repository=SimpleNamespace(),
graph_engine_layers=(),
graph_runtime_state=None,
graph_runtime_state=SimpleNamespace(),
)
queue_manager.publish_error.assert_not_called()
assert restore_workflow_run_graph.call_args.kwargs["workflow"] is workflow
assert restore_workflow_run_graph.call_args.kwargs["workflow_run_id"] == "run-id"
def test_generate_worker_handles_validation_error(self, monkeypatch: pytest.MonkeyPatch):
generator = AdvancedChatAppGenerator()
api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py
+2 -1
View File
@@ -56,6 +56,7 @@ from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType
from core.workflow.system_variables import build_system_variables
from graphon.enums import BuiltinNodeTypes
from graphon.file import FileTransferMethod, FileType
from graphon.model_runtime.entities.llm_entities import LLMUsage
from graphon.runtime import GraphRuntimeState, VariablePool
from libs.datetime_utils import naive_utc_now
from models.enums import MessageStatus
@@ -174,7 +175,7 @@ class TestAdvancedChatGenerateTaskPipeline:
variables=build_system_variables(workflow_execution_id="run-id"),
),
start_at=0.0,
total_tokens=7,
llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 7}),
node_run_steps=3,
)
api/tests/unit_tests/core/app/apps/test_base_app_generator.py
+29
View File
@@ -1,10 +1,39 @@
import logging
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
from sqlalchemy import inspect
from core.app.apps.base_app_generator import BaseAppGenerator
from graphon.variables.input_entities import VariableEntity, VariableEntityType
from models import Workflow, WorkflowRun
def test_restore_workflow_run_graph():
workflow = Workflow(graph='{"nodes": [{"id": "edited"}]}')
session = SimpleNamespace(get=Mock(return_value=SimpleNamespace(graph='{"nodes": [{"id": "paused"}]}')))
BaseAppGenerator._restore_workflow_run_graph(session=session, workflow=workflow, workflow_run_id="run-id")
session.get.assert_called_once_with(WorkflowRun, "run-id")
assert workflow.graph == '{"nodes": [{"id": "paused"}]}'
assert not inspect(workflow).attrs.graph.history.has_changes()
@pytest.mark.parametrize(
("workflow_run_id", "workflow_run"),
[(None, None), ("run-id", None), ("run-id", SimpleNamespace(graph=None))],
)
def test_restore_workflow_run_graph_requires_persisted_snapshot(workflow_run_id, workflow_run):
session = SimpleNamespace(get=Mock(return_value=workflow_run))
with pytest.raises(ValueError):
BaseAppGenerator._restore_workflow_run_graph(
session=session,
workflow=Workflow(graph="{}"),
workflow_run_id=workflow_run_id,
)
def test_validate_inputs_with_zero():
api/tests/unit_tests/core/app/apps/test_message_based_app_queue_manager.py
+24 -1
View File
@@ -6,7 +6,12 @@ from core.app.apps.base_app_queue_manager import PublishFrom
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 InvokeFrom
from core.app.entities.queue_entities import QueueErrorEvent, QueueMessageEndEvent, QueueStopEvent
from core.app.entities.queue_entities import (
QueueErrorEvent,
QueueMessageEndEvent,
QueueStopEvent,
QueueWorkflowPausedEvent,
)
class TestMessageBasedAppQueueManager:
@@ -63,3 +68,21 @@ class TestMessageBasedAppQueueManager:
manager._publish(QueueMessageEndEvent(), PublishFrom.TASK_PIPELINE)
assert manager._q.qsize() == 1
def test_publish_pause_event_stops_listener_without_aborting_execution(self):
with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis:
mock_redis.setex.return_value = True
manager = MessageBasedAppQueueManager(
task_id="t1",
user_id="u1",
invoke_from=InvokeFrom.DEBUGGER,
conversation_id="c1",
app_mode="advanced-chat",
message_id="m1",
)
manager.stop_listen = Mock()
manager._is_stopped = Mock(return_value=False)
manager._publish(QueueWorkflowPausedEvent(), PublishFrom.APPLICATION_MANAGER)
manager.stop_listen.assert_called_once_with(execution_terminal=True)
api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py
-1
View File
@@ -334,7 +334,6 @@ class TestWorkflowBasedAppRunner:
variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()),
start_at=0.0,
)
graph_runtime_state.register_paused_node("node-1")
workflow_entry = SimpleNamespace(graph_engine=SimpleNamespace(graph_runtime_state=graph_runtime_state))
emails: list[dict] = []
api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py
-3
View File
@@ -20,9 +20,6 @@ class _DummyQueueManager:
class _DummyRuntimeState:
variable_pool = object()
def get_paused_nodes(self):
return ["node-1"]
class _DummyGraphEngine:
def __init__(self):
api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py
+1
View File
@@ -130,6 +130,7 @@ def test_single_node_run_validates_target_node_config(monkeypatch: pytest.Monkey
"type": "loop",
"title": "Loop",
"loop_count": 1,
"start_node_id": "loop-start",
"break_conditions": [],
"logical_operator": "and",
},
api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py
+2 -4
View File
@@ -40,9 +40,6 @@ class _RecordingWorkflowAppRunner(WorkflowAppRunner):
class _FakeRuntimeState:
variable_pool = object()
def get_paused_nodes(self):
return ["node-pause-1"]
@pytest.fixture
def sqlite_pause_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session:
@@ -140,6 +137,7 @@ def test_graph_run_paused_event_emits_queue_pause_event(monkeypatch: pytest.Monk
"core.app.apps.workflow_app_runner.enrich_graph_pause_reasons",
lambda **_: [enriched_reason],
)
monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_email_task", MagicMock())
runner._handle_event(workflow_entry, event)
@@ -148,7 +146,7 @@ def test_graph_run_paused_event_emits_queue_pause_event(monkeypatch: pytest.Monk
assert isinstance(queue_event, QueueWorkflowPausedEvent)
assert queue_event.reasons == [enriched_reason]
assert queue_event.outputs == {"foo": "bar"}
assert queue_event.paused_nodes == ["node-pause-1"]
assert queue_event.paused_nodes == ["node-human"]
def _build_converter(*, invoke_from: InvokeFrom = InvokeFrom.SERVICE_API):
api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py
+8
View File
@@ -543,6 +543,8 @@ class TestWorkflowAppGeneratorWorker:
lambda self, *, session, workflow: workflow,
)
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)
app_config = WorkflowUIBasedAppConfig(
tenant_id="tenant",
@@ -574,6 +576,12 @@ class TestWorkflowAppGeneratorWorker:
variable_loader=SimpleNamespace(),
workflow_execution_repository=SimpleNamespace(),
workflow_node_execution_repository=SimpleNamespace(),
graph_runtime_state=SimpleNamespace(),
)
assert runner_kwargs["system_user_id"] == "session-id"
restore_workflow_run_graph.assert_called_once_with(
session=session,
workflow=workflow,
workflow_run_id="run-id",
)
api/tests/unit_tests/core/app/apps/workflow/test_app_queue_manager.py
+14 -1
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from unittest.mock import patch
from unittest.mock import Mock, patch
from core.app.apps.base_app_queue_manager import PublishFrom
from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager
@@ -41,6 +41,19 @@ class TestWorkflowAppQueueManager:
manager._publish(QueuePingEvent(), PublishFrom.TASK_PIPELINE)
def test_publish_pause_event_stops_listener_without_aborting_execution(self):
manager = WorkflowAppQueueManager(
task_id="task",
user_id="user",
invoke_from=InvokeFrom.DEBUGGER,
app_mode="workflow",
)
manager.stop_listen = Mock()
manager._publish(QueueWorkflowPausedEvent(), PublishFrom.APPLICATION_MANAGER)
manager.stop_listen.assert_called_once_with(execution_terminal=True)
def test_listener_close_aborts_unfinished_execution(self):
with (
patch("core.app.apps.base_app_queue_manager.redis_client") as redis_client,
api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py
+2 -1
View File
@@ -50,6 +50,7 @@ from core.app.entities.task_entities import (
from core.base.tts.app_generator_tts_publisher import AudioTrunk
from core.workflow.system_variables import build_system_variables, system_variables_to_mapping
from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus
from graphon.model_runtime.entities.llm_entities import LLMUsage
from graphon.runtime import GraphRuntimeState, VariablePool
from libs.datetime_utils import naive_utc_now
from models.enums import CreatorUserRole
@@ -103,7 +104,7 @@ class TestWorkflowGenerateTaskPipeline:
variables=build_system_variables(workflow_execution_id="run-id"),
),
start_at=0.0,
total_tokens=5,
llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 5}),
node_run_steps=2,
)
api/tests/unit_tests/core/app/workflow/test_persistence_layer.py
+52 -8
View File
@@ -9,7 +9,7 @@ from core.app.entities.app_invoke_entities import WorkflowAppGenerateEntity
from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer
from core.ops.ops_trace_manager import TraceTask, TraceTaskName
from core.workflow.system_variables import SystemVariableKey, build_system_variables
from graphon.entities import WorkflowNodeExecution
from graphon.entities import WorkflowNodeExecution, WorkflowStartReason
from graphon.entities.pause_reason import SchedulingPause
from graphon.enums import (
BuiltinNodeTypes,
@@ -32,6 +32,7 @@ from graphon.graph_events import (
NodeRunStartedEvent,
NodeRunSucceededEvent,
)
from graphon.model_runtime.entities.llm_entities import LLMUsage
from graphon.node_events import NodeRunResult
from graphon.runtime import GraphRuntimeState, ReadOnlyGraphRuntimeStateWrapper, VariablePool
@@ -41,6 +42,7 @@ class _RepoRecorder:
self.saved: list[object] = []
self.synchronously_saved: list[object] = []
self.saved_exec_data: list[object] = []
self.loaded: list[object] = []
def save(self, entity):
self.saved.append(entity)
@@ -51,6 +53,9 @@ class _RepoRecorder:
def save_execution_data(self, entity):
self.saved_exec_data.append(entity)
def get_by_workflow_execution(self, _workflow_execution_id):
return self.loaded
def _naive_utc_now() -> datetime:
return datetime.now(UTC).replace(tzinfo=None)
@@ -169,12 +174,45 @@ class TestWorkflowPersistenceLayer:
assert exec_repo.saved
def test_resumption_restores_container_execution_before_terminal_event(self):
layer, _, node_repo, _ = _make_layer()
started_at = _naive_utc_now()
execution = WorkflowNodeExecution(
id="loop-exec",
workflow_id="workflow-id",
workflow_execution_id="run-id",
index=4,
node_id="loop",
node_type=BuiltinNodeTypes.LOOP,
title="Loop",
status=WorkflowNodeExecutionStatus.RUNNING,
created_at=started_at,
)
node_repo.loaded = [execution]
layer.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION))
layer.on_event(
NodeRunSucceededEvent(
id=execution.id,
node_id=execution.node_id,
node_type=execution.node_type,
start_at=started_at,
node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED),
)
)
assert execution.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert layer._next_node_sequence() == 5
def test_handle_graph_run_succeeded_updates_execution(self):
layer, exec_repo, _, runtime_state = _make_layer()
layer._handle_graph_run_started()
runtime_state.total_tokens = 3
runtime_state.node_run_steps = 2
runtime_state.outputs = {"out": "v"}
usage = LLMUsage.empty_usage()
usage.total_tokens = 3
runtime_state.add_llm_usage(usage)
for _ in range(2):
runtime_state.increment_node_run_steps()
runtime_state.set_output("out", "v")
layer._handle_graph_run_succeeded(GraphRunSucceededEvent(outputs={"ok": True}))
@@ -186,8 +224,11 @@ class TestWorkflowPersistenceLayer:
def test_handle_graph_run_partial_succeeded_updates_execution(self):
layer, exec_repo, _, runtime_state = _make_layer()
layer._handle_graph_run_started()
runtime_state.total_tokens = 5
runtime_state.node_run_steps = 4
usage = LLMUsage.empty_usage()
usage.total_tokens = 5
runtime_state.add_llm_usage(usage)
for _ in range(4):
runtime_state.increment_node_run_steps()
runtime_state._graph_execution = SimpleNamespace(exceptions_count=2)
layer._handle_graph_run_partial_succeeded(
@@ -293,8 +334,11 @@ class TestWorkflowPersistenceLayer:
def test_handle_graph_run_paused_updates_outputs(self):
layer, exec_repo, _, runtime_state = _make_layer()
layer._handle_graph_run_started()
runtime_state.total_tokens = 7
runtime_state.node_run_steps = 5
usage = LLMUsage.empty_usage()
usage.total_tokens = 7
runtime_state.add_llm_usage(usage)
for _ in range(5):
runtime_state.increment_node_run_steps()
layer._handle_graph_run_paused(GraphRunPausedEvent(outputs={"pause": True}))
api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py
+54
View File
@@ -262,6 +262,60 @@ class TestCeleryWorkflowNodeExecutionRepository:
# Should return empty list since nothing in cache
assert len(result) == 0
def test_get_by_workflow_execution_loads_persisted_executions_on_cache_miss(
self, mock_session_factory, mock_account, sample_workflow_node_execution
):
repo = CeleryWorkflowNodeExecutionRepository(
session_factory=mock_session_factory,
tenant_id=RESOURCE_TENANT_ID,
user=mock_account,
app_id="test-app",
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
)
repo._sql_repository = Mock()
repo._sql_repository.get_by_workflow_execution.return_value = [sample_workflow_node_execution]
result = repo.get_by_workflow_execution(sample_workflow_node_execution.workflow_execution_id)
assert result == [sample_workflow_node_execution]
assert repo._execution_cache[sample_workflow_node_execution.id] is sample_workflow_node_execution
assert repo._workflow_execution_mapping[sample_workflow_node_execution.workflow_execution_id] == [
sample_workflow_node_execution.id
]
@patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task")
def test_get_by_workflow_execution_merges_database_and_newer_cache(
self, mock_task, mock_session_factory, mock_account, sample_workflow_node_execution
):
repo = CeleryWorkflowNodeExecutionRepository(
session_factory=mock_session_factory,
tenant_id=RESOURCE_TENANT_ID,
user=mock_account,
app_id="test-app",
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
)
persisted_current = sample_workflow_node_execution.model_copy(deep=True)
historical = sample_workflow_node_execution.model_copy(
update={
"id": str(uuid4()),
"node_execution_id": str(uuid4()),
"index": 0,
"node_id": "start",
}
)
sample_workflow_node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED
repo.save(sample_workflow_node_execution)
repo._sql_repository = Mock()
repo._sql_repository.get_by_workflow_execution.return_value = [persisted_current, historical]
result = repo.get_by_workflow_execution(
sample_workflow_node_execution.workflow_execution_id,
OrderConfig(order_by=["index"], order_direction="asc"),
)
assert [execution.id for execution in result] == [historical.id, sample_workflow_node_execution.id]
assert result[1] is sample_workflow_node_execution
@patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task")
def test_cache_operations(self, mock_task, mock_session_factory, mock_account, sample_workflow_node_execution):
"""Test cache operations work correctly."""
api/tests/unit_tests/core/tools/test_tool_manager.py
+7 -7
View File
@@ -1083,14 +1083,14 @@ def test_convert_tool_parameters_type_agent_and_workflow_branches():
variable_pool = Mock()
variable_pool.get.return_value = SimpleNamespace(value="from-variable")
variable_pool.convert_template.return_value = SimpleNamespace(text="from-template")
mixed = ToolManager._convert_tool_parameters_type(
parameters=[text_param],
variable_pool=variable_pool,
tool_configurations={"text": {"type": "mixed", "value": "Hello {{name}}"}},
typ="workflow",
)
with patch("core.tools.tool_manager.convert_template", return_value=SimpleNamespace(text="from-template")):
mixed = ToolManager._convert_tool_parameters_type(
parameters=[text_param],
variable_pool=variable_pool,
tool_configurations={"text": {"type": "mixed", "value": "Hello {{name}}"}},
typ="workflow",
)
assert mixed == {"text": "from-template"}
variable = ToolManager._convert_tool_parameters_type(
api/tests/unit_tests/core/variables/test_segment.py
+4 -3
View File
@@ -28,6 +28,7 @@ from graphon.variables.segments import (
StringSegment,
get_segment_discriminator,
)
from graphon.variables.template_resolution import convert_template
from graphon.variables.types import SegmentType
from graphon.variables.utils import (
dumps_with_segments,
@@ -98,7 +99,7 @@ def test_segment_group_to_text():
template = (
"Hello, {{#sys.user_id#}}! Your query is {{#node_id.custom_query#}}. And your key is {{#env.secret_key#}}."
)
segments_group = variable_pool.convert_template(template)
segments_group = convert_template(variable_pool, template)
assert segments_group.text == "Hello, fake-user-id! Your query is fake-user-query. And your key is fake-secret-key."
assert segments_group.log == (
@@ -112,7 +113,7 @@ def test_convert_constant_to_segment_group():
system_variables=build_system_variables(user_id="1", app_id="1", workflow_id="1"),
)
template = "Hello, world!"
segments_group = variable_pool.convert_template(template)
segments_group = convert_template(variable_pool, template)
assert segments_group.text == "Hello, world!"
assert segments_group.log == "Hello, world!"
@@ -120,7 +121,7 @@ def test_convert_constant_to_segment_group():
def test_convert_variable_to_segment_group():
variable_pool = _build_variable_pool(system_variables=build_system_variables(user_id="fake-user-id"))
template = "{{#sys.user_id#}}"
segments_group = variable_pool.convert_template(template)
segments_group = convert_template(variable_pool, template)
assert segments_group.text == "fake-user-id"
assert segments_group.log == "fake-user-id"
assert isinstance(segments_group.value[0], StringVariable)
api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py
+9 -1
View File
@@ -5,7 +5,7 @@ The factory follows the same config adaptation path as production
implementations before instantiation.
"""
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, override
from core.workflow.human_input_adapter import adapt_node_config_for_graph
from core.workflow.node_factory import DifyNodeFactory
@@ -76,6 +76,14 @@ class MockNodeFactory(DifyNodeFactory):
BuiltinNodeTypes.CODE: MockCodeNode,
}
@override
def with_runtime_state(self, graph_runtime_state: "GraphRuntimeState") -> "MockNodeFactory":
return MockNodeFactory(
graph_init_params=self.graph_init_params,
graph_runtime_state=graph_runtime_state,
mock_config=self.mock_config,
)
def create_node(self, node_config: dict[str, Any] | NodeConfigDict) -> Node:
"""
Create a node instance, using mock implementations for third-party service nodes.
api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py
-113
View File
@@ -615,69 +615,6 @@ class MockIterationNode(MockNodeMixin, IterationNode):
"""Return the version of this mock node."""
return "1"
def _create_graph_engine(self, index: int, item: Any):
"""Create a graph engine with MockNodeFactory instead of DifyNodeFactory."""
# Import dependencies
from graphon.entities import GraphInitParams
from graphon.graph import Graph
from graphon.graph_engine import GraphEngine, GraphEngineConfig
from graphon.graph_engine.command_channels import InMemoryChannel
from graphon.runtime import GraphRuntimeState
# Import our MockNodeFactory instead of DifyNodeFactory
from .test_mock_factory import MockNodeFactory
# Create GraphInitParams from node attributes
graph_init_params = GraphInitParams(
workflow_id=self.workflow_id,
graph_config=self.graph_config,
run_context=self.run_context,
call_depth=self.workflow_call_depth,
)
# Create a deep copy of the variable pool for each iteration
variable_pool_copy = self.graph_runtime_state.variable_pool.model_copy(deep=True)
# append iteration variable (item, index) to variable pool
variable_pool_copy.add([self._node_id, "index"], index)
variable_pool_copy.add([self._node_id, "item"], item)
# Create a new GraphRuntimeState for this iteration
graph_runtime_state_copy = GraphRuntimeState(
variable_pool=variable_pool_copy,
start_at=self.graph_runtime_state.start_at,
total_tokens=0,
node_run_steps=0,
)
# Create a MockNodeFactory with the same mock_config
node_factory = MockNodeFactory(
graph_init_params=graph_init_params,
graph_runtime_state=graph_runtime_state_copy,
mock_config=self.mock_config, # Pass the mock configuration
)
# Initialize the iteration graph with the mock node factory
iteration_graph = Graph.init(
graph_config=self.graph_config, node_factory=node_factory, root_node_id=self._node_data.start_node_id
)
if not iteration_graph:
from graphon.nodes.iteration.exc import IterationGraphNotFoundError
raise IterationGraphNotFoundError("iteration graph not found")
# Create a new GraphEngine for this iteration
graph_engine = GraphEngine(
workflow_id=self.workflow_id,
graph=iteration_graph,
graph_runtime_state=graph_runtime_state_copy,
command_channel=InMemoryChannel(), # Use InMemoryChannel for sub-graphs
config=GraphEngineConfig(),
)
return graph_engine
class MockLoopNode(MockNodeMixin, LoopNode):
"""Mock implementation of LoopNode that preserves mock configuration."""
@@ -687,56 +624,6 @@ class MockLoopNode(MockNodeMixin, LoopNode):
"""Return the version of this mock node."""
return "1"
def _create_graph_engine(self, start_at, root_node_id: str):
"""Create a graph engine with MockNodeFactory instead of DifyNodeFactory."""
# Import dependencies
from graphon.entities import GraphInitParams
from graphon.graph import Graph
from graphon.graph_engine import GraphEngine, GraphEngineConfig
from graphon.graph_engine.command_channels import InMemoryChannel
from graphon.runtime import GraphRuntimeState
# Import our MockNodeFactory instead of DifyNodeFactory
from .test_mock_factory import MockNodeFactory
# Create GraphInitParams from node attributes
graph_init_params = GraphInitParams(
workflow_id=self.workflow_id,
graph_config=self.graph_config,
run_context=self.run_context,
call_depth=self.workflow_call_depth,
)
# Create a new GraphRuntimeState for this iteration
graph_runtime_state_copy = GraphRuntimeState(
variable_pool=self.graph_runtime_state.variable_pool,
start_at=start_at.timestamp(),
)
# Create a MockNodeFactory with the same mock_config
node_factory = MockNodeFactory(
graph_init_params=graph_init_params,
graph_runtime_state=graph_runtime_state_copy,
mock_config=self.mock_config, # Pass the mock configuration
)
# Initialize the loop graph with the mock node factory
loop_graph = Graph.init(graph_config=self.graph_config, node_factory=node_factory, root_node_id=root_node_id)
if not loop_graph:
raise ValueError("loop graph not found")
# Create a new GraphEngine for this iteration
graph_engine = GraphEngine(
workflow_id=self.workflow_id,
graph=loop_graph,
graph_runtime_state=graph_runtime_state_copy,
command_channel=InMemoryChannel(), # Use InMemoryChannel for sub-graphs
config=GraphEngineConfig(),
)
return graph_engine
class MockTemplateTransformNode(MockNodeMixin, TemplateTransformNode):
"""Mock implementation of TemplateTransformNode for testing."""
api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py
-51
View File
@@ -51,53 +51,6 @@ from .test_mock_factory import MockNodeFactory
logger = logging.getLogger(__name__)
class _TableTestChildEngineBuilder:
def __init__(self, *, use_mock_factory: bool, mock_config: MockConfig | None) -> None:
self._use_mock_factory = use_mock_factory
self._mock_config = mock_config
def build_child_engine(
self,
*,
workflow_id: str,
graph_init_params: GraphInitParams,
parent_graph_runtime_state: GraphRuntimeState,
root_node_id: str,
variable_pool: VariablePool | None = None,
) -> GraphEngine:
child_graph_runtime_state = GraphRuntimeState(
variable_pool=variable_pool if variable_pool is not None else parent_graph_runtime_state.variable_pool,
start_at=time.perf_counter(),
execution_context=parent_graph_runtime_state.execution_context,
)
if self._use_mock_factory:
node_factory = MockNodeFactory(
graph_init_params=graph_init_params,
graph_runtime_state=child_graph_runtime_state,
mock_config=self._mock_config,
)
else:
node_factory = DifyNodeFactory(
graph_init_params=graph_init_params,
graph_runtime_state=child_graph_runtime_state,
)
graph_config = graph_init_params.graph_config
child_graph = Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id=root_node_id)
if not child_graph:
raise ValueError("child graph not found")
child_engine = GraphEngine(
workflow_id=workflow_id,
graph=child_graph,
graph_runtime_state=child_graph_runtime_state,
command_channel=InMemoryChannel(),
config=GraphEngineConfig(),
child_engine_builder=self,
)
return child_engine
@dataclass
class WorkflowTestCase:
"""Represents a single test case for table-driven testing."""
@@ -379,10 +332,6 @@ class TableTestRunner:
scale_up_threshold=self.graph_engine_scale_up_threshold,
scale_down_idle_time=self.graph_engine_scale_down_idle_time,
),
child_engine_builder=_TableTestChildEngineBuilder(
use_mock_factory=test_case.use_auto_mock,
mock_config=test_case.mock_config,
),
)
# Execute and collect events
api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py
+6 -4
View File
@@ -3,7 +3,7 @@ from datetime import UTC, datetime
from types import SimpleNamespace
from typing import cast
from unittest.mock import MagicMock, patch
from uuid import UUID
from uuid import UUID, uuid4
from agenton.compositor import CompositorSessionSnapshot
from dify_agent.layers.ask_human import AskHumanToolResult
@@ -347,7 +347,7 @@ def _node(
}
)
return DifyAgentNode(
node = DifyAgentNode(
node_id="agent-node",
data=DifyAgentNodeData.model_validate({"type": BuiltinNodeTypes.AGENT, "version": "2"}),
graph_init_params=graph_init_params,
@@ -355,7 +355,7 @@ def _node(
GraphRuntimeState,
SimpleNamespace(
variable_pool=FakeVariablePool(),
graph_execution=SimpleNamespace(node_executions={}),
graph_execution=SimpleNamespace(aborted=False),
),
),
binding_resolver=binding_resolver,
@@ -368,6 +368,8 @@ def _node(
failure_orchestrator=OutputFailureOrchestrator(),
session_store=cast(WorkflowAgentWorkspaceStore, session_store or FakeSessionStore()),
)
node.bind_execution_id(str(uuid4()))
return node
def test_extract_variable_selector_to_variable_mapping_uses_frontend_agent_task_markers():
@@ -465,7 +467,7 @@ def test_agent_node_passes_execution_id_to_session_store_and_runtime_request_bui
store = FakeSessionStore()
request_builder = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider())
node = _node(session_store=store, runtime_request_builder=request_builder)
execution_id = node.ensure_execution_id()
execution_id = node.execution_id
with patch.object(request_builder, "build", wraps=request_builder.build) as build:
list(node._run())
api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py
+3 -1
View File
@@ -73,13 +73,15 @@ def _create_human_input_node(
node_data=node_data,
file_reference_factory=_TestFileReferenceFactory(),
)
return HumanInputNode(
node = HumanInputNode(
node_id=config["id"],
data=node_data,
graph_init_params=graph_init_params,
graph_runtime_state=graph_runtime_state,
hitl_callback=callback,
)
node.bind_execution_id("00000000-0000-4000-8000-000000000001")
return node
def _build_node(
api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration_child_engine_errors.py
-95
View File
@@ -1,95 +0,0 @@
from collections.abc import Mapping
from typing import Any
import pytest
from core.workflow.system_variables import default_system_variables
from graphon.entities import GraphInitParams
from graphon.nodes.iteration.entities import IterationNodeData
from graphon.nodes.iteration.exc import IterationGraphNotFoundError
from graphon.nodes.iteration.iteration_node import IterationNode
from graphon.runtime import (
ChildEngineBuilderNotConfiguredError,
ChildGraphNotFoundError,
GraphRuntimeState,
VariablePool,
)
from tests.workflow_test_utils import build_test_graph_init_params
class _MissingGraphBuilder:
def build_child_engine(
self,
*,
workflow_id: str,
graph_init_params: GraphInitParams,
parent_graph_runtime_state: GraphRuntimeState,
root_node_id: str,
variable_pool: VariablePool | None = None,
) -> object:
raise ChildGraphNotFoundError(f"child graph root node '{root_node_id}' not found")
def _build_runtime_state() -> GraphRuntimeState:
return GraphRuntimeState(
variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables(), user_inputs={}),
start_at=0.0,
)
def _build_iteration_node(
*,
graph_config: Mapping[str, Any],
runtime_state: GraphRuntimeState,
start_node_id: str,
) -> IterationNode:
init_params = build_test_graph_init_params(graph_config=graph_config)
return IterationNode(
node_id="iteration-node",
data=IterationNodeData(
type="iteration",
title="Iteration",
iterator_selector=["start", "items"],
output_selector=["iteration-node", "output"],
start_node_id=start_node_id,
),
graph_init_params=init_params,
graph_runtime_state=runtime_state,
)
def test_graph_runtime_state_raises_specific_error_when_child_builder_is_missing():
runtime_state = _build_runtime_state()
graph_init_params = build_test_graph_init_params()
with pytest.raises(ChildEngineBuilderNotConfiguredError):
runtime_state.create_child_engine(
workflow_id="workflow",
graph_init_params=graph_init_params,
root_node_id="root",
)
def test_iteration_node_only_translates_child_graph_not_found_error():
runtime_state = _build_runtime_state()
runtime_state.bind_child_engine_builder(_MissingGraphBuilder())
node = _build_iteration_node(
graph_config={"nodes": [{"id": "present-node"}], "edges": []},
runtime_state=runtime_state,
start_node_id="missing-node",
)
with pytest.raises(IterationGraphNotFoundError):
node._create_graph_engine(index=0, item="item")
def test_iteration_node_propagates_non_graph_not_found_errors():
runtime_state = _build_runtime_state()
node = _build_iteration_node(
graph_config={"nodes": [{"id": "start-node"}], "edges": []},
runtime_state=runtime_state,
start_node_id="start-node",
)
with pytest.raises(ChildEngineBuilderNotConfiguredError):
node._create_graph_engine(index=0, item="item")
api/tests/unit_tests/core/workflow/nodes/list_operator/node_spec.py
-2
View File
@@ -1,4 +1,3 @@
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
@@ -36,7 +35,6 @@ class TestListOperatorNode:
"""Create mock GraphRuntimeState."""
mock_state = MagicMock(spec=GraphRuntimeState)
mock_variable_pool = MagicMock()
mock_variable_pool.convert_template.side_effect = lambda value: SimpleNamespace(text=value)
mock_state.variable_pool = mock_variable_pool
return mock_state
api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py
-1
View File
@@ -239,7 +239,6 @@ def test_image_link_messages_use_tool_file_id_metadata(tool_node: ToolNode):
def test_tool_node_passes_node_execution_id_when_runtime_accepts_it(tool_node: ToolNode):
runtime_handle = ToolRuntimeHandle(raw=object())
tool_node._runtime.get_runtime = MagicMock(return_value=runtime_handle)
tool_node.ensure_execution_id = MagicMock(return_value="node-execution-id")
result = tool_node._get_tool_runtime(
variable_pool=tool_node.graph_runtime_state.variable_pool,
api/tests/unit_tests/core/workflow/test_human_input_adapter.py
+3 -3
View File
@@ -18,12 +18,12 @@ from core.workflow.human_input_adapter import (
)
from graphon.enums import BuiltinNodeTypes
from graphon.nodes.base.variable_template_parser import VariableTemplateParser
from graphon.runtime import VariablePool
def test_email_delivery_config_helpers_render_and_sanitize_text() -> None:
variable_pool = SimpleNamespace(
convert_template=lambda body: SimpleNamespace(text=body.replace("{{#node.value#}}", "42"))
)
variable_pool = VariablePool()
variable_pool.add(["node", "value"], "42")
rendered = EmailDeliveryConfig.render_body_template(
body="Open {{#url#}} and use {{#node.value#}}",
api/tests/unit_tests/core/workflow/test_human_input_callback.py
+21
View File
@@ -59,6 +59,27 @@ def test_dify_hitl_callback_creates_pause_requested_for_new_form() -> None:
assert params.node_id == "node-1"
def test_dify_hitl_callback_scopes_form_to_node_execution() -> None:
repository = MagicMock(spec=HumanInputFormRepository)
repository.get_form.return_value = None
repository.create_form.return_value = SimpleNamespace(id="execution-1")
callback = DifyHITLCallback(
form_repository=repository,
node_data=HumanInputNodeData(
title="Approval",
form_content="Please approve",
user_actions=[UserActionConfig(id="approve", title="Approve")],
),
execution_id_getter=lambda: "execution-1",
)
callback(_ctx("run-1", "node-1"))
repository.get_form.assert_called_once_with("node-1", form_id="execution-1")
params: FormCreateParams = repository.create_form.call_args.args[0]
assert params.form_id == "execution-1"
def test_dify_hitl_callback_returns_completed_for_submitted_form() -> None:
repository = MagicMock(spec=HumanInputFormRepository)
repository.get_form.return_value = SimpleNamespace(
api/tests/unit_tests/core/workflow/test_node_factory.py
+13
View File
@@ -324,6 +324,19 @@ class TestDifyNodeFactoryInit:
graph_runtime_state=sentinel.graph_runtime_state,
)
def test_with_runtime_state_rebinds_factory(self):
factory = object.__new__(node_factory.DifyNodeFactory)
factory.graph_init_params = sentinel.graph_init_params
with patch.object(node_factory, "DifyNodeFactory", return_value=sentinel.factory) as factory_cls:
rebound = factory.with_runtime_state(sentinel.graph_runtime_state)
assert rebound is sentinel.factory
factory_cls.assert_called_once_with(
graph_init_params=sentinel.graph_init_params,
graph_runtime_state=sentinel.graph_runtime_state,
)
def test_init_builds_default_dependencies(self):
graph_init_params = SimpleNamespace(run_context={"context": "value"})
graph_runtime_state = sentinel.graph_runtime_state
api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py
+22 -227
View File
@@ -1,5 +1,4 @@
from collections import UserString
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import MagicMock, patch, sentinel
@@ -10,238 +9,20 @@ from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom
from core.workflow import workflow_entry
from core.workflow.system_variables import default_system_variables
from graphon.entities.base_node_data import BaseNodeData
from graphon.enums import NodeType, WorkflowNodeExecutionStatus
from graphon.enums import NodeType
from graphon.errors import WorkflowNodeRunFailedError
from graphon.file import File, FileTransferMethod, FileType
from graphon.filters import ResponseStreamFilter
from graphon.graph import Graph
from graphon.graph_events import GraphRunFailedEvent
from graphon.model_runtime.entities.llm_entities import LLMMode, LLMUsage
from graphon.node_events import NodeRunResult
from graphon.nodes import BuiltinNodeTypes
from graphon.nodes.base.node import Node
from graphon.nodes.llm.entities import ContextConfig, LLMNodeData, ModelConfig
from graphon.nodes.question_classifier.entities import QuestionClassifierNodeData
from graphon.runtime import ChildGraphNotFoundError, VariablePool
from graphon.runtime import VariablePool
from graphon.variables.variables import StringVariable
from tests.workflow_test_utils import build_test_graph_init_params, build_test_variable_pool
def _build_typed_node_config(node_type: NodeType):
return {"id": "node-id", "data": BaseNodeData(type=node_type)}
def _build_model_config(*, provider: str = "openai", model_name: str = "gpt-4o") -> ModelConfig:
return ModelConfig(provider=provider, name=model_name, mode=LLMMode.CHAT)
def _build_llm_node_data(*, provider: str = "openai", model_name: str = "gpt-4o") -> LLMNodeData:
return LLMNodeData(
type=BuiltinNodeTypes.LLM,
title="Child Model",
model=_build_model_config(provider=provider, model_name=model_name),
prompt_template=[],
context=ContextConfig(enabled=False),
)
def _build_question_classifier_node_data(
*, provider: str = "openai", model_name: str = "gpt-4o"
) -> QuestionClassifierNodeData:
return QuestionClassifierNodeData(
type=BuiltinNodeTypes.QUESTION_CLASSIFIER,
title="Child Model",
query_variable_selector=["sys", "query"],
model=_build_model_config(provider=provider, model_name=model_name),
classes=[],
)
class _FakeModelNodeMixin:
@classmethod
def version(cls) -> str:
return "1"
def post_init(self) -> None:
self.model_instance = SimpleNamespace(provider="stale-provider", model_name="stale-model")
self.usage_snapshot = LLMUsage.empty_usage()
self.usage_snapshot.total_tokens = 1
def _run(self) -> NodeRunResult:
return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED,
inputs={
"model_provider": self.node_data.model.provider,
"model_name": self.node_data.model.name,
},
llm_usage=self.usage_snapshot,
)
class _FakeLLMNode(_FakeModelNodeMixin, Node[LLMNodeData]):
node_type = BuiltinNodeTypes.LLM
class _FakeQuestionClassifierNode(_FakeModelNodeMixin, Node[QuestionClassifierNodeData]):
node_type = BuiltinNodeTypes.QUESTION_CLASSIFIER
class TestWorkflowChildEngineBuilder:
@pytest.mark.parametrize(
("graph_config", "node_id", "expected"),
[
({"nodes": [{"id": "root"}]}, "root", True),
({"nodes": [{"id": "root"}]}, "other", False),
({"nodes": "invalid"}, "root", None),
({"nodes": ["invalid"]}, "root", None),
],
)
def test_has_node_id(self, graph_config, node_id, expected):
result = workflow_entry._WorkflowChildEngineBuilder._has_node_id(graph_config, node_id)
assert result is expected
def test_build_child_engine_raises_when_root_node_is_missing(self):
builder = workflow_entry._WorkflowChildEngineBuilder(tenant_id="tenant-id")
graph_init_params = SimpleNamespace(graph_config={"nodes": []})
parent_graph_runtime_state = SimpleNamespace(
execution_context=sentinel.execution_context,
variable_pool=sentinel.variable_pool,
)
with patch.object(workflow_entry, "DifyNodeFactory", return_value=sentinel.factory):
with pytest.raises(ChildGraphNotFoundError, match="child graph root node 'missing' not found"):
builder.build_child_engine(
workflow_id="workflow-id",
graph_init_params=graph_init_params,
parent_graph_runtime_state=parent_graph_runtime_state,
root_node_id="missing",
)
def test_build_child_engine_constructs_graph_engine_with_quota_layer_only(self):
builder = workflow_entry._WorkflowChildEngineBuilder(tenant_id="tenant-id")
graph_init_params = SimpleNamespace(graph_config={"nodes": [{"id": "root"}]})
parent_graph_runtime_state = SimpleNamespace(
execution_context=sentinel.execution_context,
variable_pool=sentinel.parent_variable_pool,
)
child_graph = sentinel.child_graph
child_graph_runtime_state = sentinel.child_graph_runtime_state
child_engine = MagicMock()
with (
patch.object(workflow_entry.time, "perf_counter", return_value=123.0),
patch.object(
workflow_entry,
"GraphRuntimeState",
return_value=child_graph_runtime_state,
) as graph_runtime_state_cls,
patch.object(workflow_entry, "DifyNodeFactory", return_value=sentinel.factory) as dify_node_factory,
patch.object(workflow_entry.Graph, "init", return_value=child_graph) as graph_init,
patch.object(workflow_entry, "GraphEngine", return_value=child_engine) as graph_engine_cls,
patch.object(workflow_entry, "GraphEngineConfig", return_value=sentinel.graph_engine_config),
patch.object(workflow_entry, "InMemoryChannel", return_value=sentinel.command_channel),
patch.object(workflow_entry, "LLMQuotaLayer", return_value=sentinel.llm_quota_layer) as llm_quota_layer_cls,
):
result = builder.build_child_engine(
workflow_id="workflow-id",
graph_init_params=graph_init_params,
parent_graph_runtime_state=parent_graph_runtime_state,
root_node_id="root",
variable_pool=sentinel.child_variable_pool,
)
assert result is child_engine
graph_runtime_state_cls.assert_called_once_with(
variable_pool=sentinel.child_variable_pool,
start_at=123.0,
execution_context=sentinel.execution_context,
)
dify_node_factory.assert_called_once_with(
graph_init_params=graph_init_params,
graph_runtime_state=child_graph_runtime_state,
)
graph_init.assert_called_once_with(
graph_config={"nodes": [{"id": "root"}]},
node_factory=sentinel.factory,
root_node_id="root",
)
graph_engine_cls.assert_called_once_with(
workflow_id="workflow-id",
graph=child_graph,
graph_runtime_state=child_graph_runtime_state,
command_channel=sentinel.command_channel,
config=sentinel.graph_engine_config,
child_engine_builder=builder,
)
llm_quota_layer_cls.assert_called_once_with(tenant_id="tenant-id")
assert child_engine.layer.call_args_list == [((sentinel.llm_quota_layer,), {})]
@pytest.mark.parametrize("node_cls", [_FakeLLMNode, _FakeQuestionClassifierNode])
def test_build_child_engine_runs_llm_quota_layer_for_child_model_nodes(self, node_cls):
builder = workflow_entry._WorkflowChildEngineBuilder(tenant_id="tenant-id")
graph_init_params = build_test_graph_init_params(
graph_config={"nodes": [{"id": "root"}], "edges": []},
)
parent_graph_runtime_state = SimpleNamespace(
execution_context=nullcontext(None),
variable_pool=build_test_variable_pool(),
)
created_node: dict[str, _FakeLLMNode | _FakeQuestionClassifierNode] = {}
def build_graph(*, graph_config, node_factory, root_node_id):
_ = graph_config
node_data = _build_llm_node_data() if node_cls is _FakeLLMNode else _build_question_classifier_node_data()
node = node_cls(
node_id=root_node_id,
data=node_data,
graph_init_params=node_factory.graph_init_params,
graph_runtime_state=node_factory.graph_runtime_state,
)
created_node["node"] = node
return Graph(
nodes={root_node_id: node},
edges={},
in_edges={},
out_edges={},
root_node=node,
)
with (
patch.object(
workflow_entry,
"DifyNodeFactory",
side_effect=lambda graph_init_params, graph_runtime_state: SimpleNamespace(
graph_init_params=graph_init_params,
graph_runtime_state=graph_runtime_state,
),
),
patch.object(workflow_entry.Graph, "init", side_effect=build_graph),
patch("core.app.workflow.layers.llm_quota.ensure_llm_quota_available_for_model") as ensure_quota,
patch("core.app.workflow.layers.llm_quota.deduct_llm_quota_for_model") as deduct_quota,
):
child_engine = builder.build_child_engine(
workflow_id="workflow-id",
graph_init_params=graph_init_params,
parent_graph_runtime_state=parent_graph_runtime_state,
root_node_id="root",
)
list(child_engine.run())
node = created_node["node"]
ensure_quota.assert_called_once_with(
tenant_id="tenant-id",
provider=node.node_data.model.provider,
model=node.node_data.model.name,
)
deduct_quota.assert_called_once_with(
tenant_id="tenant-id",
provider=node.node_data.model.provider,
model=node.node_data.model.name,
usage=node.usage_snapshot,
)
def _build_minimal_workflow_entry(
monkeypatch: pytest.MonkeyPatch,
*,
@@ -249,7 +30,7 @@ def _build_minimal_workflow_entry(
) -> workflow_entry.WorkflowEntry:
"""Construct a minimal WorkflowEntry with GraphEngine construction mocked out."""
graph_engine = MagicMock()
graph_runtime_state = SimpleNamespace(execution_context=None)
graph_runtime_state = SimpleNamespace(_execution_context=None)
monkeypatch.setattr(workflow_entry, "capture_current_context", lambda: sentinel.execution_context)
monkeypatch.setattr(workflow_entry, "GraphEngine", MagicMock(return_value=graph_engine))
@@ -294,7 +75,7 @@ class TestWorkflowEntryInit:
def test_applies_debug_and_observability_layers(self):
graph_engine = MagicMock()
graph_runtime_state = SimpleNamespace(execution_context=None)
graph_runtime_state = SimpleNamespace(_execution_context=None)
debug_layer = sentinel.debug_layer
execution_limits_layer = sentinel.execution_limits_layer
llm_quota_layer = sentinel.llm_quota_layer
@@ -339,9 +120,8 @@ class TestWorkflowEntryInit:
graph_runtime_state=graph_runtime_state,
command_channel=sentinel.command_channel,
config=sentinel.graph_engine_config,
child_engine_builder=entry._child_engine_builder,
)
assert graph_runtime_state.execution_context is sentinel.execution_context
assert graph_runtime_state._execution_context is sentinel.execution_context
debug_logging_layer.assert_called_once_with(
level="DEBUG",
include_inputs=True,
@@ -443,6 +223,21 @@ class TestWorkflowEntryRun:
class TestWorkflowEntrySingleStepRun:
@pytest.mark.parametrize("node_type", [BuiltinNodeTypes.LOOP, BuiltinNodeTypes.ITERATION])
def test_rejects_container_nodes(self, node_type):
workflow = SimpleNamespace(
get_node_config_by_id=lambda _node_id: _build_typed_node_config(node_type),
)
with pytest.raises(ValueError, match="engine-backed debug endpoints"):
workflow_entry.WorkflowEntry.single_step_run(
workflow=workflow,
node_id="node-id",
user_id="user-id",
user_inputs={},
variable_pool=sentinel.variable_pool,
)
def test_preloads_constructor_variables_before_creating_memory_node(self):
class FakeLLMNode:
id = "node-id"
@@ -958,7 +753,7 @@ class TestWorkflowEntryTracing:
layer = MagicMock()
class FakeNode:
def ensure_execution_id(self):
def bind_execution_id(self, _execution_id):
return None
def run(self):
@@ -979,7 +774,7 @@ class TestWorkflowEntryTracing:
layer = MagicMock()
class FakeNode:
def ensure_execution_id(self):
def bind_execution_id(self, _execution_id):
return None
def run(self):
api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py
+25
View File
@@ -27,6 +27,7 @@ from libs.broadcast_channel.redis.sharded_channel import (
ShardedTopic,
_RedisShardedSubscription,
)
from libs.broadcast_channel.signals import SIG_CLOSE
class TestBroadcastChannel:
@@ -1239,6 +1240,30 @@ class TestRedisSubscriptionCommon:
subscription_type, _ = subscription_params
assert subscription._get_subscription_type() == subscription_type
def test_listener_ignores_close_signal_from_another_subscription(self, subscription, subscription_params):
subscription_type, _ = subscription_params
topic = f"test-{subscription_type}-topic"
message_type = "message" if subscription_type == "regular" else "smessage"
messages = iter(
[
{"type": message_type, "channel": topic, "data": SIG_CLOSE},
{"type": message_type, "channel": topic, "data": b"next-event"},
]
)
def get_message():
try:
return next(messages)
except StopIteration:
subscription._closed.set()
return None
subscription._get_message = get_message
subscription._listen()
assert subscription._queue.get_nowait() == b"next-event"
assert subscription._queue.empty()
# ==================== Lifecycle Tests ====================
def test_start_if_needed_first_call(self, subscription, subscription_params, mock_pubsub: MagicMock):
api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py
+29
View File
@@ -12,6 +12,7 @@ from libs.broadcast_channel.redis.streams_channel import (
StreamsTopic,
_StreamsSubscription,
)
from libs.broadcast_channel.signals import SIG_CLOSE
class FakeStreamsRedis:
@@ -282,6 +283,34 @@ class TestStreamsSubscription:
assert received == case.expected_messages
def test_listener_ignores_close_signal_from_another_subscription(self):
class OneShotRedis:
def __init__(self) -> None:
self._calls = 0
def xread(self, streams: dict[str, Any], block: int | None = None, count: int | None = None):
self._calls += 1
if self._calls == 1:
key = next(iter(streams))
return [
(
key,
[
("1-0", {b"data": SIG_CLOSE}),
("2-0", {b"data": b"next-event"}),
],
)
]
subscription._closed = True
return []
subscription = _StreamsSubscription(OneShotRedis(), "stream:close-signal")
subscription._listen()
assert subscription._queue.get_nowait() == b"next-event"
assert subscription._queue.get_nowait() is subscription._SENTINEL
assert subscription._queue.empty()
def test_iterator_yields_messages_until_subscription_is_closed(self, streams_channel: StreamsBroadcastChannel):
topic = streams_channel.topic("iter")
subscription = topic.subscribe()
api/tests/unit_tests/services/test_model_load_balancing_service.py
+1 -2
View File
@@ -123,12 +123,11 @@ def test_enable_disable_model_load_balancing_uses_model_type_constructor_directl
method_name: str,
expected_provider_method: str,
service: ModelLoadBalancingService,
monkeypatch: pytest.MonkeyPatch,
) -> None:
provider_configuration = _build_provider_configuration(provider_schema=_build_provider_credential_schema())
service.provider_manager.get_configurations.return_value = {"openai": provider_configuration}
getattr(service, method_name)("tenant-1", "openai", "gpt-4o-mini", "text-generation")
getattr(service, method_name)("tenant-1", "openai", "gpt-4o-mini", "llm")
getattr(provider_configuration, expected_provider_method).assert_called_once_with(
model="gpt-4o-mini", model_type=ModelType.LLM
api/tests/unit_tests/services/test_model_provider_service.py
+3 -3
View File
@@ -377,7 +377,7 @@ class TestModelProviderServiceDelegation:
{
"tenant_id": "tenant-1",
"provider": "openai",
"model_type": "text-generation",
"model_type": "llm",
"model": "gpt-4o",
"credential_id": "cred-1",
},
@@ -389,7 +389,7 @@ class TestModelProviderServiceDelegation:
{
"tenant_id": "tenant-1",
"provider": "openai",
"model_type": "text-generation",
"model_type": "llm",
"model": "gpt-4o",
"credentials": {"api_key": "x"},
"credential_name": "cred-a",
@@ -407,7 +407,7 @@ class TestModelProviderServiceDelegation:
{
"tenant_id": "tenant-1",
"provider": "openai",
"model_type": "text-generation",
"model_type": "llm",
"model": "gpt-4o",
},
"delete_custom_model",
api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py
+2 -3
View File
@@ -144,8 +144,7 @@ def _build_resumption_context(task_id: str, *, select_options: list[str] | None
runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0)
if select_options is not None:
runtime_state.variable_pool.add(("start", "options"), select_options)
runtime_state.register_paused_node("node-1")
runtime_state.outputs = {"result": "value"}
runtime_state.set_output("result", "value")
wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity)
return WorkflowResumptionContext(
generate_entity=wrapper,
@@ -250,7 +249,7 @@ def _build_resumption_context_additional(task_id: str) -> WorkflowResumptionCont
workflow_execution_id="run-1",
)
runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0)
runtime_state.outputs = {"answer": "ok"}
runtime_state.set_output("answer", "ok")
wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity)
return WorkflowResumptionContext(
generate_entity=wrapper,
api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py
+1 -1
View File
@@ -74,7 +74,7 @@ def _build_resumption_context(task_id: str) -> WorkflowResumptionContext:
workflow_execution_id="run-1",
)
runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0)
runtime_state.outputs = {"answer": "ok"}
runtime_state.set_output("answer", "ok")
wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity)
return WorkflowResumptionContext(
generate_entity=wrapper,