refactor(models): pass session into WorkflowNodeExecutionModel accessors (#41968)

api/controllers/console/app/workflow.py

@@ -77,7 +77,10 @@ from factories import file_factory, variable_factory
 from fields.base import ResponseModel
 from fields.conversation_variable_fields import WorkflowConversationVariableResponse
 from fields.member_fields import SimpleAccount
-from fields.workflow_run_fields import WorkflowRunNodeExecutionResponse
+from fields.workflow_run_fields import (
+    WorkflowRunNodeExecutionResponse,
+    node_execution_response_source,
+)
 from graphon.enums import NodeType
 from graphon.file import File
 from graphon.file import helpers as file_helpers
@@ -1260,7 +1263,7 @@ class DraftWorkflowNodeRunApi(Resource):
         )
 
         return WorkflowRunNodeExecutionResponse.model_validate(
-            workflow_node_execution, from_attributes=True
+            node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True
         ).model_dump(mode="json")
 
 
@@ -1680,7 +1683,9 @@ class DraftWorkflowNodeLastRunApi(Resource):
         )
         if node_exec is None:
             raise NotFound("last run not found")
-        return WorkflowRunNodeExecutionResponse.model_validate(node_exec, from_attributes=True).model_dump(mode="json")
+        return WorkflowRunNodeExecutionResponse.model_validate(
+            node_execution_response_source(node_exec, session=db.session()), from_attributes=True
+        ).model_dump(mode="json")
 
 
 @console_ns.route("/apps/<uuid:app_id>/workflows/draft/trigger/run")

api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py

@@ -53,6 +53,7 @@ from fields.workflow_run_fields import (
     WorkflowRunNodeExecutionListResponse,
     WorkflowRunNodeExecutionResponse,
     WorkflowRunPaginationResponse,
+    node_execution_response_source,
 )
 from graphon.model_runtime.utils.encoders import jsonable_encoder
 from libs import helper
@@ -506,7 +507,7 @@ class RagPipelineDraftNodeRunApi(Resource):
             raise ValueError("Workflow node execution not found")
 
         return WorkflowRunNodeExecutionResponse.model_validate(
-            workflow_node_execution, from_attributes=True
+            node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True
         ).model_dump(mode="json")
 
 
@@ -1023,7 +1024,9 @@ class RagPipelineWorkflowLastRunApi(Resource):
         )
         if node_exec is None:
             raise NotFound("last run not found")
-        return WorkflowRunNodeExecutionResponse.model_validate(node_exec, from_attributes=True).model_dump(mode="json")
+        return WorkflowRunNodeExecutionResponse.model_validate(
+            node_execution_response_source(node_exec, session=db.session()), from_attributes=True
+        ).model_dump(mode="json")
 
 
 @console_ns.route("/rag/pipelines/transform/datasets/<uuid:dataset_id>")
@@ -1085,7 +1088,7 @@ class RagPipelineDatasourceVariableApi(Resource):
             current_user=current_user,
         )
         return WorkflowRunNodeExecutionResponse.model_validate(
-            workflow_node_execution, from_attributes=True
+            node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True
         ).model_dump(mode="json")
 
 

api/controllers/console/snippets/snippet_workflow.py

@@ -51,6 +51,7 @@ from fields.workflow_run_fields import (
     WorkflowRunNodeExecutionListResponse,
     WorkflowRunNodeExecutionResponse,
     WorkflowRunPaginationResponse,
+    node_execution_response_source,
 )
 from graphon.graph_engine.manager import GraphEngineManager
 from libs import helper
@@ -635,7 +636,7 @@ class SnippetDraftNodeRunApi(Resource):
         )
 
         return WorkflowRunNodeExecutionResponse.model_validate(
-            workflow_node_execution, from_attributes=True
+            node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True
         ).model_dump(mode="json")
 
 
@@ -672,7 +673,9 @@ class SnippetDraftNodeLastRunApi(Resource):
         if node_exec is None:
             raise NotFound("Node last run not found")
 
-        return WorkflowRunNodeExecutionResponse.model_validate(node_exec, from_attributes=True).model_dump(mode="json")
+        return WorkflowRunNodeExecutionResponse.model_validate(
+            node_execution_response_source(node_exec, session=db.session()), from_attributes=True
+        ).model_dump(mode="json")
 
 
 @console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/iteration/nodes/<string:node_id>/run")

api/fields/workflow_run_fields.py

@@ -7,11 +7,13 @@ kept only for workflow app-log endpoints that still build legacy log models.
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import datetime
 from typing import Any
 
 from flask_restx import Namespace, fields
 from pydantic import AliasChoices, Field, field_validator
+from sqlalchemy.orm import Session
 
 from fields.base import ResponseModel
 from fields.end_user_fields import SimpleEndUser
@@ -187,5 +189,28 @@ class WorkflowRunNodeExecutionResponse(ResponseModel):
         return to_timestamp(value)
 
 
+@dataclass(frozen=True)
+class WorkflowNodeExecutionResponseSource:
+    """Expose session-backed node-execution accessors during response validation."""
+
+    node_execution: Any
+    session: Session
+
+    @property
+    def created_by_account(self) -> Any:
+        return self.node_execution.created_by_account(self.session)
+
+    @property
+    def created_by_end_user(self) -> Any:
+        return self.node_execution.created_by_end_user(self.session)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.node_execution, name)  # guard-ignore: no-new-getattr -- delegates model fields
+
+
+def node_execution_response_source(node_execution: Any, *, session: Session) -> WorkflowNodeExecutionResponseSource:
+    return WorkflowNodeExecutionResponseSource(node_execution=node_execution, session=session)
+
+
 class WorkflowRunNodeExecutionListResponse(ResponseModel):
     data: list[WorkflowRunNodeExecutionResponse]

api/models/workflow.py

@@ -1100,22 +1100,20 @@ class WorkflowNodeExecutionModel(Base):  # This model is expected to have `offlo
             )
         )
 
-    @property
-    def created_by_account(self):
+    def created_by_account(self, session: orm.Session) -> Account | None:
         created_by_role = CreatorUserRole(self.created_by_role)
         if created_by_role == CreatorUserRole.ACCOUNT:
             stmt = select(Account).where(Account.id == self.created_by)
-            return db.session.scalar(stmt)
+            return session.scalar(stmt)
         return None
 
-    @property
-    def created_by_end_user(self):
+    def created_by_end_user(self, session: orm.Session):
         from .model import EndUser
 
         created_by_role = CreatorUserRole(self.created_by_role)
         if created_by_role == CreatorUserRole.END_USER:
             stmt = select(EndUser).where(EndUser.id == self.created_by)
-            return db.session.scalar(stmt)
+            return session.scalar(stmt)
         return None
 
     @property

api/services/rag_pipeline/rag_pipeline.py

@@ -1266,7 +1266,9 @@ class RagPipelineService:
             app_id=pipeline.id,
             workflow_run_id=run_id,
         )
-        return assemble_workflow_node_execution_traces(node_executions, self._node_execution_service_repo)
+        return assemble_workflow_node_execution_traces(
+            node_executions, self._node_execution_service_repo, session=self._session
+        )
 
     @staticmethod
     def publish_customized_pipeline_template(

api/services/snippet_service.py

@@ -989,7 +989,10 @@ class SnippetService:
             workflow_run_id=workflow_run.id,
         )
 
-        return assemble_workflow_node_execution_traces(node_executions, self._node_execution_service_repo)
+        with self._session_scope() as session:
+            return assemble_workflow_node_execution_traces(
+                node_executions, self._node_execution_service_repo, session=session
+            )
 
     # --- Node Execution Operations ---
 

api/services/workflow_node_execution_trace_service.py

@@ -8,6 +8,7 @@ from enum import Enum
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict, ValidationError
+from sqlalchemy.orm import Session
 
 from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt
 from libs.helper import to_timestamp
@@ -58,23 +59,27 @@ class WorkflowNodeExecutionTrace(BaseModel):
 def assemble_workflow_node_execution_traces(
     executions: Sequence[WorkflowNodeExecutionModel],
     repository: DifyAPIWorkflowNodeExecutionRepository,
+    *,
+    session: Session,
 ) -> list[WorkflowNodeExecutionTrace]:
     """Expand valid persisted retry attempts before each terminal execution."""
     traces: list[WorkflowNodeExecutionTrace] = []
     for execution in executions:
-        traces.extend(_expand_execution(execution, repository))
+        traces.extend(_expand_execution(execution, repository, session=session))
     return traces
 
 
 def _expand_execution(
     execution: WorkflowNodeExecutionModel,
     repository: DifyAPIWorkflowNodeExecutionRepository,
+    *,
+    session: Session,
 ) -> list[WorkflowNodeExecutionTrace]:
     full_process_data = _load_full_process_data(execution, repository)
     retry_attempts = _parse_retry_attempts(execution, full_process_data)
     terminal_metadata = execution.execution_metadata_dict
-    traces = [_retry_trace(execution, attempt, terminal_metadata) for attempt in retry_attempts]
-    traces.append(_terminal_trace(execution))
+    traces = [_retry_trace(execution, attempt, terminal_metadata, session=session) for attempt in retry_attempts]
+    traces.append(_terminal_trace(execution, session=session))
     return traces
 
 
@@ -133,6 +138,8 @@ def _retry_trace(
     execution: WorkflowNodeExecutionModel,
     attempt: WorkflowNodeRetryAttempt,
     terminal_metadata: Mapping[str, Any],
+    *,
+    session: Session,
 ) -> WorkflowNodeExecutionTrace:
     truncator = VariableTruncator.default()
     inputs, inputs_truncated = truncator.truncate_variable_mapping(attempt.inputs)
@@ -163,8 +170,8 @@ def _retry_trace(
         created_at=attempt.created_at,
         created_by_role=_enum_value(execution.created_by_role),
         created_by=execution.created_by,
-        created_by_account=execution.created_by_account,
-        created_by_end_user=execution.created_by_end_user,
+        created_by_account=execution.created_by_account(session),
+        created_by_end_user=execution.created_by_end_user(session),
         finished_at=attempt.finished_at,
         inputs_truncated=inputs_truncated,
         outputs_truncated=outputs_truncated,
@@ -173,7 +180,7 @@ def _retry_trace(
     )
 
 
-def _terminal_trace(execution: WorkflowNodeExecutionModel) -> WorkflowNodeExecutionTrace:
+def _terminal_trace(execution: WorkflowNodeExecutionModel, *, session: Session) -> WorkflowNodeExecutionTrace:
     process_data = execution.process_data_dict
     if process_data is not None and RETRY_HISTORY_PROCESS_DATA_KEY in process_data:
         process_data = dict(process_data)
@@ -202,8 +209,8 @@ def _terminal_trace(execution: WorkflowNodeExecutionModel) -> WorkflowNodeExecut
         created_at=to_timestamp(execution.created_at),
         created_by_role=_enum_value(execution.created_by_role),
         created_by=execution.created_by,
-        created_by_account=execution.created_by_account,
-        created_by_end_user=execution.created_by_end_user,
+        created_by_account=execution.created_by_account(session),
+        created_by_end_user=execution.created_by_end_user(session),
         finished_at=to_timestamp(execution.finished_at),
         inputs_truncated=execution.inputs_truncated,
         outputs_truncated=execution.outputs_truncated,

api/services/workflow_run_service.py

@@ -5,6 +5,7 @@ from typing import TypedDict
 
 import contexts
 from core.workflow.nodes.human_input.pause_reason import HumanInputRequired
+from extensions.ext_database import db
 from graphon.enums import WorkflowExecutionStatus
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from machinery.context import RequestContext
@@ -202,7 +203,7 @@ class WorkflowRunService:
             app_id=app_id,
             workflow_run_id=run_id,
         )
-        return assemble_workflow_node_execution_traces(node_executions, self._node_executions)
+        return assemble_workflow_node_execution_traces(node_executions, self._node_executions, session=db.session())
 
     def get_pause_details(
         self,

api/tests/test_containers_integration_tests/models/test_workflow_node_execution_model.py

@@ -109,7 +109,7 @@ class TestWorkflowNodeExecutionModelCreatedBy:
             created_by=account.id,
         )
 
-        result = execution.created_by_account
+        result = execution.created_by_account(session=db_session_with_containers)
 
         assert result is not None
         assert result.id == account.id
@@ -126,7 +126,7 @@ class TestWorkflowNodeExecutionModelCreatedBy:
             created_by=account.id,
         )
 
-        result = execution.created_by_account
+        result = execution.created_by_account(session=db_session_with_containers)
 
         assert result is None
 
@@ -146,7 +146,7 @@ class TestWorkflowNodeExecutionModelCreatedBy:
             created_by=end_user.id,
         )
 
-        result = execution.created_by_end_user
+        result = execution.created_by_end_user(session=db_session_with_containers)
 
         assert result is not None
         assert result.id == end_user.id
@@ -165,6 +165,6 @@ class TestWorkflowNodeExecutionModelCreatedBy:
             created_by=end_user.id,
         )
 
-        result = execution.created_by_end_user
+        result = execution.created_by_end_user(session=db_session_with_containers)
 
         assert result is None

api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py

@@ -15,6 +15,7 @@ from sqlalchemy.orm import Session
 from controllers.common.errors import NotFoundError
 from controllers.console.app import workflow_run as workflow_run_module
 from extensions.ext_database import db
+from fields.workflow_run_fields import node_execution_response_source
 from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus
 from machinery.context import RequestContext
 from models import Account, App, AppMode
@@ -348,7 +349,9 @@ def test_workflow_run_node_executions_return_frontend_trace_contract(
     _account(sqlite_session)
     execution = _workflow_run_node_execution(sqlite_session)
     workflow_runs = Mock()
-    workflow_runs.get_workflow_run_node_executions.return_value = [execution]
+    workflow_runs.get_workflow_run_node_executions.return_value = [
+        node_execution_response_source(execution, session=sqlite_session)
+    ]
     _mock_application_services(monkeypatch, workflow_runs)
     monkeypatch.setattr(db, "session", sqlite_session)
     request_context = _request_context()

api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py

@@ -45,6 +45,7 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import (
     WorkflowListQuery,
     WorkflowUpdatePayload,
 )
+from fields.workflow_run_fields import node_execution_response_source
 from graphon.enums import WorkflowNodeExecutionStatus
 from libs.datetime_utils import naive_utc_now
 from models.account import Account, TenantAccountRole
@@ -802,8 +803,12 @@ class TestRagPipelineWorkflowRunNodeExecutionListApi:
         run_id = uuid4()
         node_exec = make_node_execution(workflow_run_id=str(run_id))
 
+        session_stub = MagicMock()
+        session_stub.scalar.return_value = None
         service = MagicMock()
-        service.get_rag_pipeline_workflow_run_node_executions.return_value = [node_exec]
+        service.get_rag_pipeline_workflow_run_node_executions.return_value = [
+            node_execution_response_source(node_exec, session=session_stub)
+        ]
 
         with (
             app.test_request_context("/"),

api/tests/unit_tests/models/test_workflow_node_execution_accessors.py

@@ -0,0 +1,110 @@
+"""Regression coverage for ``WorkflowNodeExecutionModel`` account accessors.
+
+Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch:
+``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user``
+looks up an EndUser only when role is END_USER. Also covers the session-carrying response
+source used by the console endpoints that validate ``WorkflowRunNodeExecutionResponse``
+with ``from_attributes=True``.
+"""
+
+from uuid import uuid4
+
+from sqlalchemy.orm import Session
+
+from fields.workflow_run_fields import node_execution_response_source
+from models.account import Account
+from models.enums import CreatorUserRole, EndUserType
+from models.model import EndUser
+from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom
+
+
+def _execution(role: CreatorUserRole, created_by: str) -> WorkflowNodeExecutionModel:
+    """Construct a WorkflowNodeExecutionModel without touching the database."""
+    return WorkflowNodeExecutionModel(
+        tenant_id="00000000-0000-0000-0000-000000000001",
+        app_id="00000000-0000-0000-0000-000000000002",
+        workflow_id=str(uuid4()),
+        triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
+        workflow_run_id=None,
+        index=1,
+        predecessor_node_id=None,
+        node_execution_id=None,
+        node_id="n1",
+        node_type="start",
+        title="Start",
+        inputs=None,
+        process_data=None,
+        outputs=None,
+        status="succeeded",
+        error=None,
+        elapsed_time=0.0,
+        execution_metadata=None,
+        created_by_role=role,
+        created_by=created_by,
+    )
+
+
+class TestCreatedByAccount:
+    def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None:
+        account = Account(name="Test Account", email="test@example.com")
+        sqlite_session.add(account)
+        sqlite_session.flush()
+        execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id)
+
+        result = execution.created_by_account(session=sqlite_session)
+
+        assert result is not None
+        assert result.id == account.id
+
+    def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None:
+        account = Account(name="Test Account", email="test@example.com")
+        sqlite_session.add(account)
+        sqlite_session.flush()
+        execution = _execution(CreatorUserRole.END_USER, created_by=account.id)
+
+        assert execution.created_by_account(session=sqlite_session) is None
+
+
+class TestCreatedByEndUser:
+    def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None:
+        end_user = EndUser(
+            tenant_id="00000000-0000-0000-0000-000000000001",
+            type=EndUserType.BROWSER,
+            session_id="session-1",
+        )
+        sqlite_session.add(end_user)
+        sqlite_session.flush()
+        execution = _execution(CreatorUserRole.END_USER, created_by=end_user.id)
+
+        result = execution.created_by_end_user(session=sqlite_session)
+
+        assert result is not None
+        assert result.id == end_user.id
+
+    def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None:
+        end_user = EndUser(
+            tenant_id="00000000-0000-0000-0000-000000000001",
+            type=EndUserType.BROWSER,
+            session_id="session-1",
+        )
+        sqlite_session.add(end_user)
+        sqlite_session.flush()
+        execution = _execution(CreatorUserRole.ACCOUNT, created_by=end_user.id)
+
+        assert execution.created_by_end_user(session=sqlite_session) is None
+
+
+class TestNodeExecutionResponseSource:
+    def test_accessors_resolve_via_wrapped_session_and_other_attributes_proxy(self, sqlite_session: Session) -> None:
+        account = Account(name="Test Account", email="test@example.com")
+        sqlite_session.add(account)
+        sqlite_session.flush()
+        execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id)
+
+        source = node_execution_response_source(execution, session=sqlite_session)
+
+        resolved = source.created_by_account
+        assert resolved is not None
+        assert resolved.id == account.id
+        assert source.created_by_end_user is None
+        assert source.node_id == "n1"

api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py

@@ -4,6 +4,7 @@ from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime
 from types import SimpleNamespace
+from unittest import mock
 
 import pytest
 from pytest_mock import MockerFixture
@@ -1498,7 +1499,7 @@ def test_get_rag_pipeline_workflow_run_node_executions_assembles_configured_repo
         app_id=pipeline.id,
         workflow_run_id="run-1",
     )
-    mock_assemble.assert_called_once_with(expected_executions, node_repo)
+    mock_assemble.assert_called_once_with(expected_executions, node_repo, session=mock.ANY)
 
 
 def test_get_recommended_plugins_returns_empty_when_no_active_plugins(

api/tests/unit_tests/services/test_snippet_service.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 from datetime import datetime
 from types import SimpleNamespace
-from unittest.mock import Mock
+from unittest.mock import ANY, Mock
 
 import pytest
 from sqlalchemy import event, select
@@ -968,6 +968,8 @@ def test_workflow_run_queries_delegate_to_repositories(monkeypatch: pytest.Monke
     )
     service._workflow_run_repo = workflow_run_repo
     service._node_execution_service_repo = node_execution_repo
+    service._session = Mock()
+    service._session_maker = Mock()
     snippet = _snippet()
     expected_traces = [SimpleNamespace(id="node-execution-1:retry:1"), SimpleNamespace(id="node-execution-1")]
     mock_assemble = Mock(return_value=expected_traces)
@@ -998,7 +1000,7 @@ def test_workflow_run_queries_delegate_to_repositories(monkeypatch: pytest.Monke
         workflow_run_id="run-1",
     )
     mock_assemble.assert_called_once_with(
-        node_execution_repo.get_executions_by_workflow_run.return_value, node_execution_repo
+        node_execution_repo.get_executions_by_workflow_run.return_value, node_execution_repo, session=ANY
     )
     node_execution_repo.get_node_last_execution.assert_called_once_with(
         tenant_id="tenant-1",

api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py

@@ -77,7 +77,7 @@ def _repository(full_process_data: dict[str, object] | None) -> DifyAPIWorkflowN
     return cast(DifyAPIWorkflowNodeExecutionRepository, repository)
 
 
-def test_assemble_expands_retry_history_before_terminal_trace() -> None:
+def test_assemble_expands_retry_history_before_terminal_trace(sqlite_session: Session) -> None:
     process_data = {
         "request": "successful-attempt",
         RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(2), _retry_attempt(1)],
@@ -85,7 +85,7 @@ def test_assemble_expands_retry_history_before_terminal_trace() -> None:
     execution = _execution(process_data)
     repository = _repository(process_data)
 
-    traces = assemble_workflow_node_execution_traces([execution], repository)
+    traces = assemble_workflow_node_execution_traces([execution], repository, session=sqlite_session)
 
     assert [trace.id for trace in traces] == ["exec-1:retry:1", "exec-1:retry:2", "exec-1"]
     assert [trace.status for trace in traces] == ["retry", "retry", "succeeded"]
@@ -104,18 +104,18 @@ def test_assemble_expands_retry_history_before_terminal_trace() -> None:
     repository.load_full_process_data.assert_called_once_with(execution)
 
 
-def test_assemble_keeps_old_execution_without_retry_history() -> None:
+def test_assemble_keeps_old_execution_without_retry_history(sqlite_session: Session) -> None:
     process_data = {"request": "terminal"}
     execution = _execution(process_data)
 
-    traces = assemble_workflow_node_execution_traces([execution], _repository(process_data))
+    traces = assemble_workflow_node_execution_traces([execution], _repository(process_data), session=sqlite_session)
 
     assert len(traces) == 1
     assert traces[0].id == "exec-1"
     assert traces[0].process_data == process_data
 
 
-def test_assemble_skips_malformed_and_duplicate_retry_attempts() -> None:
+def test_assemble_skips_malformed_and_duplicate_retry_attempts(sqlite_session: Session) -> None:
     process_data = {
         RETRY_HISTORY_PROCESS_DATA_KEY: [
             _retry_attempt(2),
@@ -125,35 +125,41 @@ def test_assemble_skips_malformed_and_duplicate_retry_attempts() -> None:
         ]
     }
 
-    traces = assemble_workflow_node_execution_traces([_execution(process_data)], _repository(process_data))
+    traces = assemble_workflow_node_execution_traces(
+        [_execution(process_data)], _repository(process_data), session=sqlite_session
+    )
 
     assert [trace.retry_index for trace in traces[:-1]] == [1, 2]
     assert traces[0].error == "attempt 1 failed"
 
 
-def test_assemble_truncates_retry_attempt_fields() -> None:
+def test_assemble_truncates_retry_attempt_fields(sqlite_session: Session) -> None:
     process_data = {RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(1, outputs={"body": "x" * 2_000_000})]}
 
-    traces = assemble_workflow_node_execution_traces([_execution(process_data)], _repository(process_data))
+    traces = assemble_workflow_node_execution_traces(
+        [_execution(process_data)], _repository(process_data), session=sqlite_session
+    )
 
     assert traces[0].outputs_truncated is True
     assert traces[-1].id == "exec-1"
 
 
-def test_assemble_falls_back_to_inline_process_data_when_loader_fails() -> None:
+def test_assemble_falls_back_to_inline_process_data_when_loader_fails(sqlite_session: Session) -> None:
     process_data = {RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(1)]}
     execution = _execution(process_data)
     repository = _repository(None)
     repository.load_full_process_data.side_effect = OSError("storage unavailable")
 
-    traces = assemble_workflow_node_execution_traces([execution], repository)
+    traces = assemble_workflow_node_execution_traces([execution], repository, session=sqlite_session)
 
     assert [trace.id for trace in traces] == ["exec-1:retry:1", "exec-1"]
 
 
-def test_virtual_trace_validates_through_node_execution_response() -> None:
+def test_virtual_trace_validates_through_node_execution_response(sqlite_session: Session) -> None:
     process_data = {RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(1)]}
-    traces = assemble_workflow_node_execution_traces([_execution(process_data)], _repository(process_data))
+    traces = assemble_workflow_node_execution_traces(
+        [_execution(process_data)], _repository(process_data), session=sqlite_session
+    )
 
     response = WorkflowRunNodeExecutionListResponse.model_validate({"data": traces}, from_attributes=True)
 

api/tests/unit_tests/services/test_workflow_run_service.py

@@ -1,7 +1,7 @@
 """Unit tests for the Console workflow-run application service."""
 
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import ANY, MagicMock
 
 import pytest
 
@@ -213,4 +213,4 @@ class TestWorkflowRunServiceQueries:
             app_id="app-1",
             workflow_run_id="run-1",
         )
-        mock_assemble.assert_called_once_with(expected_executions, node_executions)
+        mock_assemble.assert_called_once_with(expected_executions, node_executions, session=ANY)