refactor(models): pass session into WorkflowDraftVariable value decoding (#43112)

api/controllers/console/app/workflow_draft_variable.py

@@ -7,12 +7,13 @@ from uuid import UUID
 from flask import Response
 from flask_restx import Resource
 from pydantic import BaseModel, ConfigDict, Field, model_validator
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 from controllers.common.errors import InvalidArgumentError, NotFoundError
 from controllers.common.fields import SimpleResultResponse
 from controllers.common.rbac import PlainApp, RBACCheck
 from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
+from controllers.common.session import with_session
 from controllers.console import console_ns
 from controllers.console.app.error import (
     DraftWorkflowNotExist,
@@ -41,6 +42,8 @@ from fields.workflow_draft_variable_fields import (
     WorkflowDraftVariableListWithoutValueResponse,
     WorkflowDraftVariableResponse,
     WorkflowDraftVariableWithoutValueResponse,
+    draft_variable_list_response_source,
+    draft_variable_response_source,
 )
 from graphon.variables.types import SegmentType
 from libs.external_api import error_responses
@@ -248,12 +251,12 @@ class WorkflowVariableCollectionApi(Resource):
     @console_ns.doc(description="Delete all draft workflow variables")
     @console_ns.response(204, "Workflow variables deleted successfully")
     @_api_prerequisite
-    def delete(self, current_user: Account, app_model: App):
+    @with_session
+    def delete(self, session: Session, current_user: Account, app_model: App):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
         draft_var_srv.delete_user_workflow_variables(app_model.id, user_id=current_user.id)
-        db.session.commit()
         return Response("", 204)
 
 
@@ -286,25 +289,27 @@ class NodeVariableCollectionApi(Resource):
     )
     @_api_prerequisite
     @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp()))
-    def get(self, current_user: Account, app_model: App, node_id: str):
+    @with_session(write=False)
+    def get(self, session: Session, current_user: Account, app_model: App, node_id: str):
         validate_node_id(node_id)
-        with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
-            draft_var_srv = WorkflowDraftVariableService(
-                session=session,
-            )
-            node_vars = draft_var_srv.list_node_variables(app_model.id, node_id, user_id=current_user.id)
-
-        return dump_response(WorkflowDraftVariableListResponse, node_vars)
+        draft_var_srv = WorkflowDraftVariableService(
+            session=session,
+        )
+        node_vars = draft_var_srv.list_node_variables(app_model.id, node_id, user_id=current_user.id)
+        return dump_response(
+            WorkflowDraftVariableListResponse,
+            draft_variable_list_response_source(node_vars, session=session),
+        )
 
     @console_ns.doc("delete_node_variables")
     @console_ns.doc(description="Delete all variables for a specific node")
     @console_ns.response(204, "Node variables deleted successfully")
     @_api_prerequisite
-    def delete(self, current_user: Account, app_model: App, node_id: str):
+    @with_session
+    def delete(self, session: Session, current_user: Account, app_model: App, node_id: str):
         validate_node_id(node_id)
-        srv = WorkflowDraftVariableService(db.session())
+        srv = WorkflowDraftVariableService(session)
         srv.delete_node_variables(app_model.id, node_id, user_id=current_user.id)
-        db.session.commit()
         return Response("", 204)
 
 
@@ -324,9 +329,10 @@ class VariableApi(Resource):
     @console_ns.response(404, "Variable not found")
     @_api_prerequisite
     @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp()))
-    def get(self, current_user: Account, app_model: App, variable_id: UUID):
+    @with_session(write=False)
+    def get(self, session: Session, current_user: Account, app_model: App, variable_id: UUID):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
         variable_id_str = str(variable_id)
         variable = ensure_variable_access(
@@ -335,7 +341,7 @@ class VariableApi(Resource):
             variable_id=variable_id_str,
             current_user_id=current_user.id,
         )
-        return dump_response(WorkflowDraftVariableResponse, variable)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session))
 
     @console_ns.doc("update_variable")
     @console_ns.doc(description="Update a workflow variable")
@@ -348,8 +354,10 @@ class VariableApi(Resource):
     @console_ns.response(404, "Variable not found")
     @_api_prerequisite
     @model_validate(WorkflowDraftVariableUpdatePayload)
+    @with_session
     def patch(
         self,
+        session: Session,
         req_data: WorkflowDraftVariableUpdatePayload,
         current_user: Account,
         app_model: App,
@@ -377,7 +385,7 @@ class VariableApi(Resource):
         #     }
 
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
 
         variable_id_str = str(variable_id)
@@ -391,7 +399,9 @@ class VariableApi(Resource):
         new_name = req_data.name
         raw_value = req_data.value
         if new_name is None and raw_value is None:
-            return dump_response(WorkflowDraftVariableResponse, variable)
+            return dump_response(
+                WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)
+            )
 
         new_value = None
         if raw_value is not None:
@@ -418,17 +428,17 @@ class VariableApi(Resource):
                     pass
             new_value = build_segment_with_type(variable.value_type, raw_value)
         draft_var_srv.update_variable(variable, name=new_name, value=new_value)
-        db.session.commit()
-        return dump_response(WorkflowDraftVariableResponse, variable)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session))
 
     @console_ns.doc("delete_variable")
     @console_ns.doc(description="Delete a workflow variable")
     @console_ns.response(204, "Variable deleted successfully")
     @console_ns.response(404, "Variable not found")
     @_api_prerequisite
-    def delete(self, current_user: Account, app_model: App, variable_id: UUID):
+    @with_session
+    def delete(self, session: Session, current_user: Account, app_model: App, variable_id: UUID):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
         variable_id_str = str(variable_id)
         variable = ensure_variable_access(
@@ -438,7 +448,6 @@ class VariableApi(Resource):
             current_user_id=current_user.id,
         )
         draft_var_srv.delete_variable(variable)
-        db.session.commit()
         return Response("", 204)
 
 
@@ -455,13 +464,14 @@ class VariableResetApi(Resource):
     @console_ns.response(204, "Variable reset (no content)")
     @console_ns.response(404, "Variable not found")
     @_api_prerequisite
-    def put(self, current_user: Account, app_model: App, variable_id: UUID):
+    @with_session
+    def put(self, session: Session, current_user: Account, app_model: App, variable_id: UUID):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
 
         workflow_srv = WorkflowService()
-        draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session())
+        draft_workflow = workflow_srv.get_draft_workflow(app_model, session=session)
         if draft_workflow is None:
             raise NotFoundError(
                 f"Draft workflow not found, app_id={app_model.id}",
@@ -475,27 +485,27 @@ class VariableResetApi(Resource):
         )
 
         resetted = draft_var_srv.reset_variable(draft_workflow, variable)
-        db.session.commit()
         if resetted is None:
             return Response("", 204)
-        return dump_response(WorkflowDraftVariableResponse, resetted)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session))
 
 
-def _get_variable_list(app_model: App, node_id: str, current_user_id: str) -> WorkflowDraftVariableList:
-    with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
-        draft_var_srv = WorkflowDraftVariableService(
-            session=session,
+def _get_variable_list(
+    app_model: App, node_id: str, current_user_id: str, *, session: Session
+) -> WorkflowDraftVariableList:
+    draft_var_srv = WorkflowDraftVariableService(
+        session=session,
+    )
+    if node_id == CONVERSATION_VARIABLE_NODE_ID:
+        draft_vars = draft_var_srv.list_conversation_variables(app_model.id, user_id=current_user_id)
+    elif node_id == SYSTEM_VARIABLE_NODE_ID:
+        draft_vars = draft_var_srv.list_system_variables(app_model.id, user_id=current_user_id)
+    else:
+        draft_vars = draft_var_srv.list_node_variables(
+            app_id=app_model.id,
+            node_id=node_id,
+            user_id=current_user_id,
         )
-        if node_id == CONVERSATION_VARIABLE_NODE_ID:
-            draft_vars = draft_var_srv.list_conversation_variables(app_model.id, user_id=current_user_id)
-        elif node_id == SYSTEM_VARIABLE_NODE_ID:
-            draft_vars = draft_var_srv.list_system_variables(app_model.id, user_id=current_user_id)
-        else:
-            draft_vars = draft_var_srv.list_node_variables(
-                app_id=app_model.id,
-                node_id=node_id,
-                user_id=current_user_id,
-            )
     return draft_vars
 
 
@@ -512,19 +522,22 @@ class ConversationVariableCollectionApi(Resource):
     @console_ns.response(404, "Draft workflow not found")
     @_api_prerequisite
     @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp()))
-    def get(self, current_user: Account, app_model: App):
+    @with_session
+    def get(self, session: Session, current_user: Account, app_model: App):
         # NOTE(QuantumGhost): Prefill conversation variables into the draft variables table
         # so their IDs can be returned to the caller.
         workflow_srv = WorkflowService()
-        draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session())
+        draft_workflow = workflow_srv.get_draft_workflow(app_model, session=session)
         if draft_workflow is None:
             raise NotFoundError(description=f"draft workflow not found, id={app_model.id}")
-        draft_var_srv = WorkflowDraftVariableService(db.session())
+        draft_var_srv = WorkflowDraftVariableService(session)
         draft_var_srv.prefill_conversation_variable_default_values(draft_workflow, user_id=current_user.id)
-        db.session.commit()
         return dump_response(
             WorkflowDraftVariableListResponse,
-            _get_variable_list(app_model, CONVERSATION_VARIABLE_NODE_ID, current_user.id),
+            draft_variable_list_response_source(
+                _get_variable_list(app_model, CONVERSATION_VARIABLE_NODE_ID, current_user.id, session=session),
+                session=session,
+            ),
         )
 
     @console_ns.expect(console_ns.models[ConversationVariableUpdatePayload.__name__])
@@ -576,10 +589,14 @@ class SystemVariableCollectionApi(Resource):
     )
     @_api_prerequisite
     @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp()))
-    def get(self, current_user: Account, app_model: App):
+    @with_session(write=False)
+    def get(self, session: Session, current_user: Account, app_model: App):
         return dump_response(
             WorkflowDraftVariableListResponse,
-            _get_variable_list(app_model, SYSTEM_VARIABLE_NODE_ID, current_user.id),
+            draft_variable_list_response_source(
+                _get_variable_list(app_model, SYSTEM_VARIABLE_NODE_ID, current_user.id, session=session),
+                session=session,
+            ),
         )
 
 

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

@@ -7,11 +7,12 @@ from uuid import UUID
 from flask import Response
 from flask_restx import Resource
 from pydantic import BaseModel, Field
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 from controllers.common.errors import InvalidArgumentError, NotFoundError
 from controllers.common.rbac import DatasetByPipeline, RBACCheck
 from controllers.common.schema import query_params_from_model, register_schema_models
+from controllers.common.session import with_session
 from controllers.console import console_ns
 from controllers.console.app.error import (
     DraftWorkflowNotExist,
@@ -40,6 +41,8 @@ from fields.workflow_draft_variable_fields import (
     WorkflowDraftVariableListResponse,
     WorkflowDraftVariableListWithoutValueResponse,
     WorkflowDraftVariableResponse,
+    draft_variable_list_response_source,
+    draft_variable_response_source,
 )
 from graphon.variables.types import SegmentType
 from libs.external_api import error_responses
@@ -132,12 +135,12 @@ class RagPipelineVariableCollectionApi(Resource):
 
     @console_ns.response(204, "Workflow variables deleted successfully")
     @_api_prerequisite
-    def delete(self, current_user: Account, pipeline: Pipeline):
+    @with_session
+    def delete(self, session: Session, current_user: Account, pipeline: Pipeline):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
         draft_var_srv.delete_user_workflow_variables(pipeline.id, user_id=current_user.id)
-        db.session.commit()
         return Response("", 204)
 
 
@@ -167,23 +170,25 @@ class RagPipelineNodeVariableCollectionApi(Resource):
         console_ns.models[WorkflowDraftVariableListResponse.__name__],
     )
     @_api_prerequisite
-    def get(self, current_user: Account, pipeline: Pipeline, node_id: str):
+    @with_session(write=False)
+    def get(self, session: Session, current_user: Account, pipeline: Pipeline, node_id: str):
         validate_node_id(node_id)
-        with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
-            draft_var_srv = WorkflowDraftVariableService(
-                session=session,
-            )
-            node_vars = draft_var_srv.list_node_variables(pipeline.id, node_id, user_id=current_user.id)
-
-        return dump_response(WorkflowDraftVariableListResponse, node_vars)
+        draft_var_srv = WorkflowDraftVariableService(
+            session=session,
+        )
+        node_vars = draft_var_srv.list_node_variables(pipeline.id, node_id, user_id=current_user.id)
+        return dump_response(
+            WorkflowDraftVariableListResponse,
+            draft_variable_list_response_source(node_vars, session=session),
+        )
 
     @console_ns.response(204, "Node variables deleted successfully")
     @_api_prerequisite
-    def delete(self, current_user: Account, pipeline: Pipeline, node_id: str):
+    @with_session
+    def delete(self, session: Session, current_user: Account, pipeline: Pipeline, node_id: str):
         validate_node_id(node_id)
-        srv = WorkflowDraftVariableService(db.session())
+        srv = WorkflowDraftVariableService(session)
         srv.delete_node_variables(pipeline.id, node_id, user_id=current_user.id)
-        db.session.commit()
         return Response("", 204)
 
 
@@ -198,9 +203,10 @@ class RagPipelineVariableApi(Resource):
         console_ns.models[WorkflowDraftVariableResponse.__name__],
     )
     @_api_prerequisite
-    def get(self, current_user: Account, pipeline: Pipeline, variable_id: UUID):
+    @with_session(write=False)
+    def get(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
         variable_id_str = str(variable_id)
         variable = ensure_variable_access(
@@ -209,7 +215,7 @@ class RagPipelineVariableApi(Resource):
             variable_id=variable_id_str,
             current_user_id=current_user.id,
         )
-        return dump_response(WorkflowDraftVariableResponse, variable)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session))
 
     @console_ns.response(
         200,
@@ -219,8 +225,10 @@ class RagPipelineVariableApi(Resource):
     @_api_prerequisite
     @console_ns.expect(console_ns.models[WorkflowDraftVariablePatchPayload.__name__])
     @model_validate(WorkflowDraftVariablePatchPayload)
+    @with_session
     def patch(
         self,
+        session: Session,
         req_data: WorkflowDraftVariablePatchPayload,
         current_user: Account,
         pipeline: Pipeline,
@@ -248,7 +256,7 @@ class RagPipelineVariableApi(Resource):
         #     }
 
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
         args = req_data.model_dump(exclude_none=True)
 
@@ -263,7 +271,9 @@ class RagPipelineVariableApi(Resource):
         new_name = args.get(self._PATCH_NAME_FIELD, None)
         raw_value = args.get(self._PATCH_VALUE_FIELD, None)
         if new_name is None and raw_value is None:
-            return dump_response(WorkflowDraftVariableResponse, variable)
+            return dump_response(
+                WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)
+            )
 
         new_value = None
         if raw_value is not None:
@@ -290,14 +300,14 @@ class RagPipelineVariableApi(Resource):
                     pass
             new_value = build_segment_with_type(variable.value_type, raw_value)
         draft_var_srv.update_variable(variable, name=new_name, value=new_value)
-        db.session.commit()
-        return dump_response(WorkflowDraftVariableResponse, variable)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session))
 
     @console_ns.response(204, "Variable deleted successfully")
     @_api_prerequisite
-    def delete(self, current_user: Account, pipeline: Pipeline, variable_id: UUID):
+    @with_session
+    def delete(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
         variable_id_str = str(variable_id)
         variable = ensure_variable_access(
@@ -307,7 +317,6 @@ class RagPipelineVariableApi(Resource):
             current_user_id=current_user.id,
         )
         draft_var_srv.delete_variable(variable)
-        db.session.commit()
         return Response("", 204)
 
 
@@ -320,12 +329,13 @@ class RagPipelineVariableResetApi(Resource):
     )
     @console_ns.response(204, "Variable reset (no content)")
     @_api_prerequisite
-    def put(self, current_user: Account, pipeline: Pipeline, variable_id: UUID):
+    @with_session
+    def put(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID):
         draft_var_srv = WorkflowDraftVariableService(
-            session=db.session(),
+            session=session,
         )
 
-        rag_pipeline_service = RagPipelineService(db.session())
+        rag_pipeline_service = RagPipelineService(session)
         draft_workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline)
         if draft_workflow is None:
             raise NotFoundError(
@@ -340,23 +350,23 @@ class RagPipelineVariableResetApi(Resource):
         )
 
         resetted = draft_var_srv.reset_variable(draft_workflow, variable)
-        db.session.commit()
         if resetted is None:
             return Response("", 204)
-        return dump_response(WorkflowDraftVariableResponse, resetted)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session))
 
 
-def _get_variable_list(pipeline: Pipeline, node_id: str, current_user_id: str) -> WorkflowDraftVariableList:
-    with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
-        draft_var_srv = WorkflowDraftVariableService(
-            session=session,
-        )
-        if node_id == CONVERSATION_VARIABLE_NODE_ID:
-            draft_vars = draft_var_srv.list_conversation_variables(pipeline.id, user_id=current_user_id)
-        elif node_id == SYSTEM_VARIABLE_NODE_ID:
-            draft_vars = draft_var_srv.list_system_variables(pipeline.id, user_id=current_user_id)
-        else:
-            draft_vars = draft_var_srv.list_node_variables(app_id=pipeline.id, node_id=node_id, user_id=current_user_id)
+def _get_variable_list(
+    pipeline: Pipeline, node_id: str, current_user_id: str, *, session: Session
+) -> WorkflowDraftVariableList:
+    draft_var_srv = WorkflowDraftVariableService(
+        session=session,
+    )
+    if node_id == CONVERSATION_VARIABLE_NODE_ID:
+        draft_vars = draft_var_srv.list_conversation_variables(pipeline.id, user_id=current_user_id)
+    elif node_id == SYSTEM_VARIABLE_NODE_ID:
+        draft_vars = draft_var_srv.list_system_variables(pipeline.id, user_id=current_user_id)
+    else:
+        draft_vars = draft_var_srv.list_node_variables(app_id=pipeline.id, node_id=node_id, user_id=current_user_id)
     return draft_vars
 
 
@@ -368,10 +378,14 @@ class RagPipelineSystemVariableCollectionApi(Resource):
         console_ns.models[WorkflowDraftVariableListResponse.__name__],
     )
     @_api_prerequisite
-    def get(self, current_user: Account, pipeline: Pipeline):
+    @with_session(write=False)
+    def get(self, session: Session, current_user: Account, pipeline: Pipeline):
         return dump_response(
             WorkflowDraftVariableListResponse,
-            _get_variable_list(pipeline, SYSTEM_VARIABLE_NODE_ID, current_user.id),
+            draft_variable_list_response_source(
+                _get_variable_list(pipeline, SYSTEM_VARIABLE_NODE_ID, current_user.id, session=session),
+                session=session,
+            ),
         )
 
 

api/controllers/console/snippets/snippet_workflow_draft_variable.py

@@ -48,6 +48,8 @@ from fields.workflow_draft_variable_fields import (
     WorkflowDraftVariableListResponse,
     WorkflowDraftVariableListWithoutValueResponse,
     WorkflowDraftVariableResponse,
+    draft_variable_list_response_source,
+    draft_variable_response_source,
 )
 from graphon.variables.types import SegmentType
 from libs.external_api import error_responses
@@ -160,7 +162,9 @@ class SnippetNodeVariableCollectionApi(Resource):
         draft_var_srv = WorkflowDraftVariableService(session=session)
         node_vars = draft_var_srv.list_node_variables(snippet.id, node_id, user_id=current_user.id)
 
-        return dump_response(WorkflowDraftVariableListResponse, node_vars)
+        return dump_response(
+            WorkflowDraftVariableListResponse, draft_variable_list_response_source(node_vars, session=session)
+        )
 
     @console_ns.doc("delete_snippet_node_variables")
     @console_ns.doc(description="Delete all variables for a specific node (snippet draft workflow)")
@@ -197,7 +201,7 @@ class SnippetVariableApi(Resource):
             current_user_id=current_user.id,
         )
         _ensure_snippet_draft_variable_row_allowed(variable=variable, variable_id=variable_id)
-        return dump_response(WorkflowDraftVariableResponse, variable)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session))
 
     @console_ns.doc("update_snippet_workflow_variable")
     @console_ns.doc(description="Update a draft workflow variable (snippet scope)")
@@ -232,7 +236,9 @@ class SnippetVariableApi(Resource):
         new_name = req_data.name
         raw_value = req_data.value
         if new_name is None and raw_value is None:
-            return dump_response(WorkflowDraftVariableResponse, variable)
+            return dump_response(
+                WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)
+            )
 
         new_value = None
         if raw_value is not None:
@@ -256,7 +262,7 @@ class SnippetVariableApi(Resource):
                 )
             new_value = build_segment_with_type(variable.value_type, raw_value)
         draft_var_srv.update_variable(variable, name=new_name, value=new_value)
-        return dump_response(WorkflowDraftVariableResponse, variable)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session))
 
     @console_ns.doc("delete_snippet_workflow_variable")
     @console_ns.doc(description="Delete a draft workflow variable (snippet scope)")
@@ -311,7 +317,7 @@ class SnippetVariableResetApi(Resource):
         resetted = draft_var_srv.reset_variable(draft_workflow, variable)
         if resetted is None:
             return Response("", 204)
-        return dump_response(WorkflowDraftVariableResponse, resetted)
+        return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session))
 
 
 @console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/conversation-variables")

api/fields/workflow_draft_variable_fields.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 from typing import Any, override
 
 from pydantic import model_validator
+from sqlalchemy.orm import Session
 
-from fields.base import ResponseModel
+from fields.base import ResponseModel, SessionResponseSource
 from graphon.file import helpers as file_helpers
 from graphon.variables.segment_group import SegmentGroup
 from graphon.variables.segments import ArrayFileSegment, FileSegment, Segment
@@ -25,7 +26,27 @@ def _convert_values_to_json_serializable_object(value: Segment) -> RecursiveJSON
             return value.value
 
 
-def _serialize_var_value(variable: WorkflowDraftVariable) -> RecursiveJSONValue:
+class WorkflowDraftVariableResponseSource(SessionResponseSource[WorkflowDraftVariable]):
+    """Expose the session-backed value decoding during response validation."""
+
+    def get_value(self) -> Segment:
+        return self._source.get_value(session=self._session)
+
+
+def draft_variable_response_source(
+    variable: WorkflowDraftVariable, *, session: Session
+) -> WorkflowDraftVariableResponseSource:
+    return WorkflowDraftVariableResponseSource(variable, session=session)
+
+
+def draft_variable_list_response_source(variable_list: Any, *, session: Session) -> dict[str, Any]:
+    """Wrap each variable in a list payload so value decoding resolves via the session."""
+    return {
+        "items": [draft_variable_response_source(variable, session=session) for variable in variable_list.variables],
+    }
+
+
+def _serialize_var_value(variable: WorkflowDraftVariableResponseSource) -> RecursiveJSONValue:
     value = variable.get_value()
     # Create a copy to avoid mutating the model's cached deserialized value.
     value = value.model_copy(deep=True)
@@ -105,12 +126,18 @@ class WorkflowDraftVariableResponse(WorkflowDraftVariableWithoutValueResponse):
     @classmethod
     @override
     def _from_workflow_draft_variable(cls, value: Any) -> Any:
-        if isinstance(value, WorkflowDraftVariable):
+        if isinstance(value, WorkflowDraftVariableResponseSource):
             return {
-                **_serialize_without_value(value),
+                **_serialize_without_value(value._source),
                 "value": _serialize_var_value(value),
-                "full_content": _serialize_full_content(value),
+                "full_content": _serialize_full_content(value._source),
             }
+        if isinstance(value, WorkflowDraftVariable):
+            # Decoding the value needs a database session; a bare model would fall through to
+            # from_attributes and expose the raw serialized JSON string instead of the decoded value.
+            raise TypeError(
+                "WorkflowDraftVariableResponse requires draft_variable_response_source(variable, session=...)"
+            )
         return value
 
 

api/models/workflow.py

@@ -22,7 +22,7 @@ from sqlalchemy import (
     orm,
     select,
 )
-from sqlalchemy.orm import Mapped, mapped_column
+from sqlalchemy.orm import Mapped, Session, mapped_column
 
 from constants.trigger import TRIGGER_PLUGIN_NODE_TYPE, is_trigger_node_type
 from core.workflow.environment_variables import load_environment_variables
@@ -74,7 +74,6 @@ from libs import helper
 
 from .account import Account
 from .base import DefaultFieldsDCMixin, TypeBase
-from .engine import db
 from .enums import CreatorUserRole, DraftVariableType, ExecutionOffLoadType, WorkflowRunTriggeredFrom
 
 # UploadFile and workflow execution models use TypeBase; importing the class object keeps
@@ -92,10 +91,10 @@ SerializedWorkflowValue = dict[str, Any]
 SerializedWorkflowVariables = dict[str, SerializedWorkflowValue]
 
 
-def _resolve_workflow_app_tenant_id(app_id: str) -> str:
+def _resolve_workflow_app_tenant_id(session: Session, app_id: str) -> str:
     from .model import App
 
-    tenant_id = db.session.scalar(select(App.tenant_id).where(App.id == app_id))
+    tenant_id = session.scalar(select(App.tenant_id).where(App.id == app_id))
     if not tenant_id:
         raise ValueError(f"Unable to resolve tenant_id for app {app_id}")
     return tenant_id
@@ -1775,11 +1774,11 @@ class WorkflowDraftVariable(TypeBase):
     def _set_selector(self, value: list[str]):
         self.selector = json.dumps(value)
 
-    def _loads_value(self) -> Segment:
+    def _loads_value(self, *, session: Session) -> Segment:
         value = json.loads(self.value)
-        return self.build_segment_from_serialized_value(self.value_type, value)
+        return self.build_segment_from_serialized_value(self.value_type, value, session=session)
 
-    def _rebuild_file_types(self, value: Any):
+    def _rebuild_file_types(self, value: Any, *, session: Session):
         # NOTE(QuantumGhost): Temporary workaround for structured data handling.
         # By this point, `output` has been converted to dict by
         # `WorkflowEntry.handle_special_values`, so we need to
@@ -1794,7 +1793,7 @@ class WorkflowDraftVariable(TypeBase):
             case dict():
                 if not maybe_file_object(value):
                     return cast(Any, value)
-                tenant_id = _resolve_workflow_app_tenant_id(self.app_id)
+                tenant_id = _resolve_workflow_app_tenant_id(session, self.app_id)
                 return build_file_from_stored_mapping(
                     file_mapping=cast(dict[str, Any], value),
                     tenant_id=tenant_id,
@@ -1804,7 +1803,7 @@ class WorkflowDraftVariable(TypeBase):
                 first: Any = value_list[0]
                 if not maybe_file_object(first):
                     return cast(Any, value)
-                tenant_id = _resolve_workflow_app_tenant_id(self.app_id)
+                tenant_id = _resolve_workflow_app_tenant_id(session, self.app_id)
                 file_list: list[File] = []
                 for item in value_list:
                     file_list.append(
@@ -1817,7 +1816,9 @@ class WorkflowDraftVariable(TypeBase):
             case _:
                 return cast(Any, value)
 
-    def build_segment_from_serialized_value(self, segment_type: SegmentType, value: Any) -> Segment:
+    def build_segment_from_serialized_value(
+        self, segment_type: SegmentType, value: Any, *, session: Session
+    ) -> Segment:
         # Persisted draft variable rows may contain historical file payloads.
         # Rebuild them through the file factory so tenant ownership, signed URLs,
         # and storage-backed metadata come from canonical records instead of the
@@ -1828,14 +1829,14 @@ class WorkflowDraftVariable(TypeBase):
                     case File():
                         return build_segment_with_type(segment_type, value)
                     case dict():
-                        file = self._rebuild_file_types(value)
+                        file = self._rebuild_file_types(value, session=session)
                         return build_segment_with_type(segment_type, file)
                     case _:
                         raise TypeMismatchError(f"expected dict or File for FileSegment, got {type(value)}")
             case SegmentType.ARRAY_FILE:
                 if not isinstance(value, list):
                     raise TypeMismatchError(f"expected list for ArrayFileSegment, got {type(value)}")
-                file_list = self._rebuild_file_types(value)
+                file_list = self._rebuild_file_types(value, session=session)
                 return build_segment_with_type(segment_type=segment_type, value=file_list)
             case _:
                 return build_segment_with_type(segment_type=segment_type, value=value)
@@ -1890,7 +1891,7 @@ class WorkflowDraftVariable(TypeBase):
             case _:
                 return build_segment_with_type(segment_type=segment_type, value=value)
 
-    def get_value(self) -> Segment:
+    def get_value(self, *, session: Session) -> Segment:
         """Decode the serialized value into its corresponding `Segment` object.
 
         This method caches the result, so repeated calls will return the same
@@ -1908,7 +1909,7 @@ class WorkflowDraftVariable(TypeBase):
 
         if self.__value is not None:
             return self.__value
-        value = self._loads_value()
+        value = self._loads_value(session=session)
         self.__value = value
         return value
 

api/services/workflow_draft_variable_service.py

@@ -114,21 +114,37 @@ class DraftVarLoader(VariableLoader):
         # Map each selector (as a tuple via `_selector_to_tuple`) to its corresponding variable instance.
         variable_by_selector: dict[tuple[str, str], VariableBase] = {}
 
+        files: list[File] = []
+        offloaded_draft_vars: list[WorkflowDraftVariable] = []
         with Session(bind=self._engine, expire_on_commit=False) as session:
             srv = WorkflowDraftVariableService(session)
             draft_vars = srv.get_draft_variables_by_selectors(self._app_id, selectors, user_id=self._user_id)
 
-        # Important:
-        files: list[File] = []
-        # FileSegment and ArrayFileSegment are not subject to offloading, so their values
-        # can be safely accessed before any offloading logic is applied.
-        for draft_var in draft_vars:
-            value = draft_var.get_value()
-            match value:
-                case FileSegment():
-                    files.append(value.value)
-                case ArrayFileSegment():
-                    files.extend(value.value)
+            for draft_var in draft_vars:
+                if draft_var.is_truncated():
+                    offloaded_draft_vars.append(draft_var)
+                    continue
+
+                # FileSegment and ArrayFileSegment are not subject to offloading, so their values
+                # can be safely accessed before any offloading logic is applied.
+                segment = draft_var.get_value(session=session)
+                match segment:
+                    case FileSegment():
+                        files.append(segment.value)
+                    case ArrayFileSegment():
+                        files.extend(segment.value)
+                variable = segment_to_variable(
+                    segment=segment,
+                    selector=draft_var.get_selector(),
+                    variable_id=draft_var.id,
+                    name=draft_var.name,
+                    description=draft_var.description,
+                )
+                selector_tuple = self._selector_to_tuple(variable.selector)
+                variable_by_selector[selector_tuple] = variable
+
+        # The files above are the same instances held by the variables, so loading their
+        # storage keys here is visible through `variable_by_selector` as well.
         with Session(bind=self._engine) as session:
             storage_key_loader = StorageKeyLoader(
                 session,
@@ -137,23 +153,6 @@ class DraftVarLoader(VariableLoader):
             )
             storage_key_loader.load_storage_keys(files)
 
-        offloaded_draft_vars = []
-        for draft_var in draft_vars:
-            if draft_var.is_truncated():
-                offloaded_draft_vars.append(draft_var)
-                continue
-
-            segment = draft_var.get_value()
-            variable = segment_to_variable(
-                segment=segment,
-                selector=draft_var.get_selector(),
-                variable_id=draft_var.id,
-                name=draft_var.name,
-                description=draft_var.description,
-            )
-            selector_tuple = self._selector_to_tuple(variable.selector)
-            variable_by_selector[selector_tuple] = variable
-
         # Load offloaded variables using multithreading.
         # This approach reduces loading time by querying external systems concurrently.
         with ThreadPoolExecutor(max_workers=10) as executor:
@@ -188,7 +187,11 @@ class DraftVarLoader(VariableLoader):
             return (draft_var.node_id, draft_var.name), variable
 
         deserialized = json.loads(content)
-        segment = draft_var.build_segment_from_serialized_value(variable_file.value_type, deserialized)
+        # Runs on a worker thread, so it opens its own session instead of sharing one.
+        with Session(bind=self._engine, expire_on_commit=False) as session:
+            segment = draft_var.build_segment_from_serialized_value(
+                variable_file.value_type, deserialized, session=session
+            )
         variable = segment_to_variable(
             segment=segment,
             selector=draft_var.get_selector(),
@@ -573,7 +576,7 @@ class WorkflowDraftVariableService:
         )
         if draft_var is None:
             return None
-        segment = draft_var.get_value()
+        segment = draft_var.get_value(session=self._session)
         if not isinstance(segment, StringSegment):
             logger.warning(
                 "sys.conversation_id variable is not a string: app_id=%s, user_id=%s, id=%s",

api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py

@@ -288,6 +288,9 @@ def test_variable_api_patch_success(
     assert payload["id"] == variable.id
     assert payload["name"] == "renamed_var"
 
+    # The handler commits through its own decorator-injected session, so expire this
+    # session's cached instance to read the persisted row instead of the identity map.
+    db_session_with_containers.expire_all()
     refreshed = db_session_with_containers.scalar(
         select(WorkflowDraftVariable).where(WorkflowDraftVariable.id == variable.id)
     )

api/tests/e2e_tests/containers/services/test_workflow_draft_variable_service.py

@@ -213,7 +213,7 @@ class TestWorkflowDraftVariableService:
         assert retrieved_variable.id == variable.id
         assert retrieved_variable.name == "test_var"
         assert retrieved_variable.app_id == app.id
-        assert retrieved_variable.get_value().value == test_value.value
+        assert retrieved_variable.get_value(session=db_session_with_containers).value == test_value.value
 
     def test_get_variable_not_found(self, db_session_with_containers: Session, mock_external_service_dependencies):
         """
@@ -274,11 +274,11 @@ class TestWorkflowDraftVariableService:
         assert "var3" in var_names
         for var in retrieved_variables:
             if var.name == "var1":
-                assert var.get_value().value == var1_value.value
+                assert var.get_value(session=db_session_with_containers).value == var1_value.value
             elif var.name == "var2":
-                assert var.get_value().value == var2_value.value
+                assert var.get_value(session=db_session_with_containers).value == var2_value.value
             elif var.name == "var3":
-                assert var.get_value().value == var3_value.value
+                assert var.get_value(session=db_session_with_containers).value == var3_value.value
 
     def test_list_variables_without_values_success(
         self, db_session_with_containers: Session, mock_external_service_dependencies
@@ -431,12 +431,12 @@ class TestWorkflowDraftVariableService:
         service = WorkflowDraftVariableService(db_session_with_containers)
         updated_variable = service.update_variable(variable, name="new_name", value=new_value)
         assert updated_variable.name == "new_name"
-        assert updated_variable.get_value().value == new_value.value
+        assert updated_variable.get_value(session=db_session_with_containers).value == new_value.value
         assert updated_variable.last_edited_at is not None
 
         db_session_with_containers.refresh(variable)
         assert variable.name == "new_name"
-        assert variable.get_value().value == new_value.value
+        assert variable.get_value(session=db_session_with_containers).value == new_value.value
         assert variable.last_edited_at is not None
 
     def test_update_variable_not_editable(
@@ -509,10 +509,10 @@ class TestWorkflowDraftVariableService:
         service = WorkflowDraftVariableService(db_session_with_containers)
         reset_variable = service.reset_variable(workflow, variable)
         assert reset_variable is not None
-        assert reset_variable.get_value().value == "default_value"
+        assert reset_variable.get_value(session=db_session_with_containers).value == "default_value"
         assert reset_variable.last_edited_at is None
         db_session_with_containers.refresh(variable)
-        assert variable.get_value().value == "default_value"
+        assert variable.get_value(session=db_session_with_containers).value == "default_value"
         assert variable.last_edited_at is None
 
     def test_delete_variable_success(self, db_session_with_containers: Session, mock_external_service_dependencies):

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

@@ -0,0 +1,219 @@
+"""Direct-call tests for the app draft-variable handlers converted to `@with_session`."""
+
+import json
+from dataclasses import dataclass
+from inspect import unwrap
+
+import pytest
+from flask import Flask
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session
+
+from controllers.console.app import workflow_draft_variable as module
+from graphon.variables.segments import StringSegment
+from models.account import Account, AccountStatus
+from models.model import App, AppMode
+from models.workflow import Workflow, WorkflowDraftVariable
+
+_APP_ID = "app-1"
+_USER_ID = "user-1"
+
+
+def _make_account() -> Account:
+    account = Account(
+        name="tester",
+        email="tester@example.com",
+        status=AccountStatus.ACTIVE,
+    )
+    account.id = _USER_ID  # type: ignore[assignment]
+    return account
+
+
+def _make_app_model() -> App:
+    return App(
+        id=_APP_ID,
+        tenant_id="tenant-1",
+        name="App",
+        mode=AppMode.ADVANCED_CHAT,
+        enable_site=True,
+        enable_api=True,
+    )
+
+
+def _make_draft_workflow() -> Workflow:
+    workflow = Workflow.new(
+        tenant_id="tenant-1",
+        app_id=_APP_ID,
+        type="workflow",
+        version=Workflow.VERSION_DRAFT,
+        graph=json.dumps({"nodes": [], "edges": []}),
+        features="{}",
+        created_by=_USER_ID,
+        environment_variables=[],
+        conversation_variables=[],
+        rag_pipeline_variables=[],
+    )
+    workflow.id = "workflow-1"
+    return workflow
+
+
+def _make_node_variable(variable_id: str, *, node_id: str = "llm-1") -> WorkflowDraftVariable:
+    variable = WorkflowDraftVariable.new_node_variable(
+        app_id=_APP_ID,
+        user_id=_USER_ID,
+        node_id=node_id,
+        name=variable_id,
+        value=StringSegment(value=f"value-{variable_id}"),
+        node_execution_id="execution-1",
+    )
+    variable.id = variable_id
+    return variable
+
+
+def _make_system_variable(variable_id: str, *, name: str) -> WorkflowDraftVariable:
+    variable = WorkflowDraftVariable.new_sys_variable(
+        app_id=_APP_ID,
+        user_id=_USER_ID,
+        name=name,
+        value=StringSegment(value=f"value-{name}"),
+        node_execution_id="execution-1",
+    )
+    variable.id = variable_id
+    return variable
+
+
+@dataclass
+class _Database:
+    engine: Engine
+
+
+@pytest.fixture
+def app(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Flask:
+    app = Flask("test_workflow_draft_variable_handlers")
+    monkeypatch.setattr("services.workflow_service.db", _Database(sqlite_engine))
+    app.config["TESTING"] = True
+    return app
+
+
+def test_node_variable_collection_get_returns_decoded_values(app: Flask, sqlite_session: Session) -> None:
+    variable = _make_node_variable("var-1")
+    sqlite_session.add(variable)
+    sqlite_session.commit()
+    api = module.NodeVariableCollectionApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), node_id="llm-1")
+
+    assert [item["id"] for item in result["items"]] == ["var-1"]
+    assert result["items"][0]["value"] == "value-var-1"
+
+
+def test_variable_get_returns_decoded_value(app: Flask, sqlite_session: Session) -> None:
+    variable = _make_node_variable("var-1")
+    sqlite_session.add(variable)
+    sqlite_session.commit()
+    api = module.VariableApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), variable_id="var-1")
+
+    assert result["id"] == "var-1"
+    assert result["value"] == "value-var-1"
+
+
+def test_variable_patch_without_changes_returns_variable(app: Flask, sqlite_session: Session) -> None:
+    variable = _make_node_variable("var-1")
+    sqlite_session.add(variable)
+    sqlite_session.commit()
+    api = module.VariableApi()
+    handler = unwrap(api.patch)
+
+    with app.test_request_context("/", method="PATCH", json={}):
+        result = handler(
+            api,
+            sqlite_session,
+            module.WorkflowDraftVariableUpdatePayload(),
+            _make_account(),
+            app_model=_make_app_model(),
+            variable_id="var-1",
+        )
+
+    assert result["id"] == "var-1"
+    assert result["value"] == "value-var-1"
+
+
+def test_variable_patch_renames_variable(app: Flask, sqlite_session: Session) -> None:
+    variable = _make_node_variable("var-1")
+    sqlite_session.add(variable)
+    sqlite_session.commit()
+    api = module.VariableApi()
+    handler = unwrap(api.patch)
+
+    with app.test_request_context("/", method="PATCH", json={"name": "renamed"}):
+        result = handler(
+            api,
+            sqlite_session,
+            module.WorkflowDraftVariableUpdatePayload(name="renamed"),
+            _make_account(),
+            app_model=_make_app_model(),
+            variable_id="var-1",
+        )
+
+    assert result["name"] == "renamed"
+    stored = sqlite_session.get(WorkflowDraftVariable, "var-1")
+    assert stored is not None
+    assert stored.name == "renamed"
+
+
+def test_variable_reset_returns_default_value(app: Flask, sqlite_session: Session) -> None:
+    workflow = _make_draft_workflow()
+    variable = _make_node_variable("var-1")
+    variable.node_execution_id = None
+    sqlite_session.add_all([workflow, variable])
+    sqlite_session.commit()
+    api = module.VariableResetApi()
+    handler = unwrap(api.put)
+
+    with app.test_request_context("/", method="PUT"):
+        response = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), variable_id="var-1")
+
+    # A node variable without an associated execution has no default to restore, so the
+    # handler answers 204 after removing the row.
+    assert response.status_code == 204
+
+
+def test_conversation_variable_collection_get_lists_conversation_variables(app: Flask, sqlite_session: Session) -> None:
+    workflow = _make_draft_workflow()
+    conv_var = WorkflowDraftVariable.new_conversation_variable(
+        app_id=_APP_ID,
+        name="conv_var",
+        value=StringSegment(value="conv-value"),
+    )
+    conv_var.id = "conv-1"
+    conv_var.user_id = _USER_ID
+    sqlite_session.add_all([workflow, conv_var])
+    sqlite_session.commit()
+    api = module.ConversationVariableCollectionApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model())
+
+    assert [item["id"] for item in result["items"]] == ["conv-1"]
+    assert result["items"][0]["value"] == "conv-value"
+
+
+def test_system_variable_collection_get_lists_system_variables(app: Flask, sqlite_session: Session) -> None:
+    sys_var = _make_system_variable("sys-1", name="query")
+    sqlite_session.add(sys_var)
+    sqlite_session.commit()
+    api = module.SystemVariableCollectionApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model())
+
+    assert [item["id"] for item in result["items"]] == ["sys-1"]
+    assert result["items"][0]["value"] == "value-query"

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

@@ -178,7 +178,7 @@ class TestRagPipelineVariableCollectionApi:
             patch("controllers.console.datasets.rag_pipeline.rag_pipeline_draft_variable.db", sqlite_db),
             patch("controllers.console.datasets.rag_pipeline.rag_pipeline_draft_variable.WorkflowDraftVariableService"),
         ):
-            result = method(api, editor_user, pipeline)
+            result = method(api, sqlite_db.session(), editor_user, pipeline)
 
         assert isinstance(result, Response)
         assert result.status_code == 204
@@ -206,12 +206,12 @@ class TestRagPipelineNodeVariableCollectionApi:
                 return_value=srv,
             ),
         ):
-            result = method(api, editor_user, pipeline, "node1")
+            result = method(api, sqlite_db.session(), editor_user, pipeline, "node1")
 
         assert result == {"items": []}
         srv.list_node_variables.assert_called_once_with("p1", "node1", user_id="account-1")
 
-    def test_get_node_variables_invalid_node(self, app: Flask, editor_user):
+    def test_get_node_variables_invalid_node(self, app: Flask, sqlite_db, editor_user):
         api = RagPipelineNodeVariableCollectionApi()
         method = unwrap(api.get)
 
@@ -219,7 +219,7 @@ class TestRagPipelineNodeVariableCollectionApi:
             app.test_request_context("/"),
         ):
             with pytest.raises(InvalidArgumentError):
-                method(api, editor_user, _pipeline(), SYSTEM_VARIABLE_NODE_ID)
+                method(api, sqlite_db.session(), editor_user, _pipeline(), SYSTEM_VARIABLE_NODE_ID)
 
 
 class TestRagPipelineVariableApi:
@@ -239,7 +239,7 @@ class TestRagPipelineVariableApi:
             ),
         ):
             with pytest.raises(NotFoundError):
-                method(api, editor_user, _pipeline(), "v1")
+                method(api, sqlite_db.session(), editor_user, _pipeline(), "v1")
 
     def test_patch_variable_invalid_file_payload(self, app: Flask, sqlite_db, editor_user):
         api = RagPipelineVariableApi()
@@ -263,7 +263,14 @@ class TestRagPipelineVariableApi:
             ),
         ):
             with pytest.raises(InvalidArgumentError):
-                method(api, WorkflowDraftVariablePatchPayload.model_validate(payload), editor_user, pipeline, "v1")
+                method(
+                    api,
+                    sqlite_db.session(),
+                    WorkflowDraftVariablePatchPayload.model_validate(payload),
+                    editor_user,
+                    pipeline,
+                    "v1",
+                )
 
     def test_delete_variable_success(self, app: Flask, sqlite_db, editor_user):
         api = RagPipelineVariableApi()
@@ -283,7 +290,7 @@ class TestRagPipelineVariableApi:
                 return_value=srv,
             ),
         ):
-            result = method(api, editor_user, pipeline, "v1")
+            result = method(api, sqlite_db.session(), editor_user, pipeline, "v1")
 
         assert result.status_code == 204
 
@@ -318,9 +325,9 @@ def test_direct_variable_access_rejects_different_user(
     rag_service = MagicMock()
     rag_service.get_draft_workflow.return_value = _workflow(pipeline)
     if payload is not None:
-        call_args = (api, payload, editor_user, pipeline, "v1")
+        call_args = (api, sqlite_db.session(), payload, editor_user, pipeline, "v1")
     else:
-        call_args = (api, editor_user, pipeline, "v1")
+        call_args = (api, sqlite_db.session(), editor_user, pipeline, "v1")
 
     with (
         app.test_request_context("/"),
@@ -376,7 +383,7 @@ class TestRagPipelineVariableResetApi:
                 return_value={"id": "v1"},
             ),
         ):
-            result = method(api, editor_user, pipeline, "v1")
+            result = method(api, sqlite_db.session(), editor_user, pipeline, "v1")
 
         assert result == {"id": "v1"}
 
@@ -403,7 +410,7 @@ class TestSystemAndEnvironmentVariablesApi:
                 return_value=srv,
             ),
         ):
-            result = method(api, editor_user, pipeline)
+            result = method(api, sqlite_db.session(), editor_user, pipeline)
 
         assert result == {"items": []}
         srv.list_system_variables.assert_called_once_with("p1", user_id="account-1")

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

@@ -0,0 +1,166 @@
+"""Direct-call tests for the rag-pipeline draft-variable handlers converted to `@with_session`."""
+
+import json
+from dataclasses import dataclass
+from inspect import unwrap
+
+import pytest
+from flask import Flask
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session
+
+from controllers.console.datasets.rag_pipeline import rag_pipeline_draft_variable as module
+from graphon.variables.segments import StringSegment
+from models.account import Account, AccountStatus
+from models.dataset import Pipeline
+from models.workflow import Workflow, WorkflowDraftVariable
+
+_PIPELINE_ID = "pipeline-1"
+_USER_ID = "user-1"
+
+
+def _make_account() -> Account:
+    account = Account(
+        name="tester",
+        email="tester@example.com",
+        status=AccountStatus.ACTIVE,
+    )
+    account.id = _USER_ID  # type: ignore[assignment]
+    return account
+
+
+def _make_pipeline() -> Pipeline:
+    pipeline = Pipeline(
+        tenant_id="tenant-1",
+        name="Pipeline",
+    )
+    pipeline.id = _PIPELINE_ID
+    return pipeline
+
+
+def _make_draft_workflow() -> Workflow:
+    workflow = Workflow.new(
+        tenant_id="tenant-1",
+        app_id=_PIPELINE_ID,
+        type="rag-pipeline",
+        version=Workflow.VERSION_DRAFT,
+        graph=json.dumps({"nodes": [], "edges": []}),
+        features="{}",
+        created_by=_USER_ID,
+        environment_variables=[],
+        conversation_variables=[],
+        rag_pipeline_variables=[],
+    )
+    workflow.id = "workflow-1"
+    return workflow
+
+
+def _make_node_variable(variable_id: str, *, node_id: str = "extract-1") -> WorkflowDraftVariable:
+    variable = WorkflowDraftVariable.new_node_variable(
+        app_id=_PIPELINE_ID,
+        user_id=_USER_ID,
+        node_id=node_id,
+        name=variable_id,
+        value=StringSegment(value=f"value-{variable_id}"),
+        node_execution_id="execution-1",
+    )
+    variable.id = variable_id
+    return variable
+
+
+@dataclass
+class _Database:
+    engine: Engine
+
+
+@pytest.fixture
+def app(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Flask:
+    app = Flask("test_rag_pipeline_draft_variable_handlers")
+    monkeypatch.setattr("services.rag_pipeline.rag_pipeline.db", _Database(sqlite_engine))
+    app.config["TESTING"] = True
+    return app
+
+
+def test_node_variable_collection_get_returns_decoded_values(app: Flask, sqlite_session: Session) -> None:
+    variable = _make_node_variable("var-1")
+    sqlite_session.add(variable)
+    sqlite_session.commit()
+    api = module.RagPipelineNodeVariableCollectionApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), node_id="extract-1")
+
+    assert [item["id"] for item in result["items"]] == ["var-1"]
+    assert result["items"][0]["value"] == "value-var-1"
+
+
+def test_variable_get_returns_decoded_value(app: Flask, sqlite_session: Session) -> None:
+    variable = _make_node_variable("var-1")
+    sqlite_session.add(variable)
+    sqlite_session.commit()
+    api = module.RagPipelineVariableApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), variable_id="var-1")
+
+    assert result["id"] == "var-1"
+    assert result["value"] == "value-var-1"
+
+
+def test_variable_patch_without_changes_returns_variable(app: Flask, sqlite_session: Session) -> None:
+    variable = _make_node_variable("var-1")
+    sqlite_session.add(variable)
+    sqlite_session.commit()
+    api = module.RagPipelineVariableApi()
+    handler = unwrap(api.patch)
+
+    with app.test_request_context("/", method="PATCH", json={}):
+        result = handler(
+            api,
+            sqlite_session,
+            module.WorkflowDraftVariablePatchPayload(),
+            _make_account(),
+            pipeline=_make_pipeline(),
+            variable_id="var-1",
+        )
+
+    assert result["id"] == "var-1"
+    assert result["value"] == "value-var-1"
+
+
+def test_variable_reset_removes_variable_without_default(app: Flask, sqlite_session: Session) -> None:
+    workflow = _make_draft_workflow()
+    variable = _make_node_variable("var-1")
+    variable.node_execution_id = None
+    sqlite_session.add_all([workflow, variable])
+    sqlite_session.commit()
+    api = module.RagPipelineVariableResetApi()
+    handler = unwrap(api.put)
+
+    with app.test_request_context("/", method="PUT"):
+        response = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), variable_id="var-1")
+
+    assert response.status_code == 204
+
+
+def test_system_variable_collection_get_lists_system_variables(app: Flask, sqlite_session: Session) -> None:
+    sys_var = WorkflowDraftVariable.new_sys_variable(
+        app_id=_PIPELINE_ID,
+        user_id=_USER_ID,
+        name="query",
+        value=StringSegment(value="sys-value"),
+        node_execution_id="execution-1",
+    )
+    sys_var.id = "sys-1"
+    sqlite_session.add(sys_var)
+    sqlite_session.commit()
+    api = module.RagPipelineSystemVariableCollectionApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline())
+
+    assert [item["id"] for item in result["items"]] == ["sys-1"]
+    assert result["items"][0]["value"] == "sys-value"

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

@@ -18,7 +18,6 @@ from sqlalchemy import Engine
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 from werkzeug.exceptions import BadRequest, Forbidden, HTTPException, NotFound
 
-import models.workflow as workflow_models
 import services
 import services.rag_pipeline.rag_pipeline as rag_pipeline_service
 from controllers.console.app.error import DraftWorkflowNotExist, DraftWorkflowNotSync
@@ -989,7 +988,6 @@ class TestRagPipelineWorkflowApisCases:
         database_session = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False))
         database = SQLiteDatabase(engine=sqlite_engine, session=database_session)
         monkeypatch.setattr(workflow_controller, "db", database)
-        monkeypatch.setattr(workflow_models, "db", database)
         monkeypatch.setattr(rag_pipeline_service, "db", database)
 
         with database_session() as session:

api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py

@@ -241,6 +241,46 @@ def test_variable_patch_returns_persisted_variable_without_committing_when_no_ch
     assert result["id"] == variable.id
 
 
+def test_variable_get_returns_decoded_persisted_value(app: Flask, sqlite_session: Session) -> None:
+    """Given a persisted snippet variable, return its decoded value through the caller's session."""
+    variable = _make_node_variable("var-1")
+    _persist_variables(sqlite_session, variable)
+    sqlite_session.expunge_all()
+    api = module.SnippetVariableApi()
+    handler = unwrap(api.get)
+
+    with app.test_request_context("/"):
+        result = handler(api, sqlite_session, _make_account(), snippet=_make_snippet(), variable_id="var-1")
+
+    assert result["id"] == "var-1"
+    assert result["value"] == "value-var-1"
+
+
+def test_variable_patch_persists_and_returns_decoded_updated_value(app: Flask, sqlite_session: Session) -> None:
+    """Given a persisted snippet variable, update its value and return the decoded result without losing the row."""
+    variable = _make_node_variable("var-1")
+    _persist_variables(sqlite_session, variable)
+    sqlite_session.expunge_all()
+    api = module.SnippetVariableApi()
+    handler = unwrap(api.patch)
+
+    with app.test_request_context("/", method="PATCH", json={"value": "updated-value"}):
+        result = handler(
+            api,
+            module.WorkflowDraftVariableUpdatePayload(value="updated-value"),
+            sqlite_session,
+            _make_account(),
+            snippet=_make_snippet(),
+            variable_id="var-1",
+        )
+
+    assert result["value"] == "updated-value"
+    sqlite_session.expire_all()
+    stored = sqlite_session.get(WorkflowDraftVariable, "var-1")
+    assert stored is not None
+    assert stored.get_value(session=sqlite_session).value == "updated-value"
+
+
 def test_variable_delete_deletes_persisted_variable(
     app: Flask,
     sqlite_session: Session,

api/tests/unit_tests/controllers/service_api/app/test_workflow.py

@@ -160,7 +160,6 @@ def _bind_sqlite_database(
     """Bind controller- and model-owned database access to the test engine."""
     database = SimpleNamespace(engine=sqlite_engine, session=sqlite_session)
     monkeypatch.setattr(sys.modules["controllers.service_api.app.workflow"], "db", database)
-    monkeypatch.setattr(sys.modules["models.workflow"], "db", database)
 
 
 def _persist_workflow_log(

api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py

@@ -4,6 +4,7 @@ from typing import Any, NamedTuple
 from unittest.mock import patch
 
 import pytest
+from sqlalchemy.orm import Session
 
 from core.workflow.variable_prefixes import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID
 from factories.variable_factory import build_segment
@@ -13,6 +14,8 @@ from fields.workflow_draft_variable_fields import (
     WorkflowDraftVariableResponse,
     WorkflowDraftVariableWithoutValueResponse,
     _serialize_full_content,
+    draft_variable_list_response_source,
+    draft_variable_response_source,
 )
 from graphon.variables.types import SegmentType
 from libs.datetime_utils import naive_utc_now
@@ -82,7 +85,7 @@ class TestWorkflowDraftVariableFields:
         with pytest.raises(AssertionError):
             _serialize_full_content(draft_var)
 
-    def test_conversation_variable(self) -> None:
+    def test_conversation_variable(self, sqlite_session: Session) -> None:
         conv_var = WorkflowDraftVariable.new_conversation_variable(
             app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1)
         )
@@ -108,9 +111,14 @@ class TestWorkflowDraftVariableFields:
         expected_with_value = expected_without_value.copy()
         expected_with_value["value"] = 1
         expected_with_value["full_content"] = None
-        assert dump_response(WorkflowDraftVariableResponse, conv_var) == expected_with_value
+        assert (
+            dump_response(
+                WorkflowDraftVariableResponse, draft_variable_response_source(conv_var, session=sqlite_session)
+            )
+            == expected_with_value
+        )
 
-    def test_create_sys_variable(self) -> None:
+    def test_create_sys_variable(self, sqlite_session: Session) -> None:
         sys_var = WorkflowDraftVariable.new_sys_variable(
             app_id=_TEST_APP_ID,
             name="sys_var",
@@ -140,9 +148,14 @@ class TestWorkflowDraftVariableFields:
         expected_with_value = expected_without_value.copy()
         expected_with_value["value"] = "a"
         expected_with_value["full_content"] = None
-        assert dump_response(WorkflowDraftVariableResponse, sys_var) == expected_with_value
+        assert (
+            dump_response(
+                WorkflowDraftVariableResponse, draft_variable_response_source(sys_var, session=sqlite_session)
+            )
+            == expected_with_value
+        )
 
-    def test_node_variable(self) -> None:
+    def test_node_variable(self, sqlite_session: Session) -> None:
         node_var = WorkflowDraftVariable.new_node_variable(
             app_id=_TEST_APP_ID,
             node_id="test_node",
@@ -173,9 +186,14 @@ class TestWorkflowDraftVariableFields:
         expected_with_value = expected_without_value.copy()
         expected_with_value["value"] = [1, "a"]
         expected_with_value["full_content"] = None
-        assert dump_response(WorkflowDraftVariableResponse, node_var) == expected_with_value
+        assert (
+            dump_response(
+                WorkflowDraftVariableResponse, draft_variable_response_source(node_var, session=sqlite_session)
+            )
+            == expected_with_value
+        )
 
-    def test_node_variable_with_file(self) -> None:
+    def test_node_variable_with_file(self, sqlite_session: Session) -> None:
         node_var = WorkflowDraftVariable.new_node_variable(
             app_id=_TEST_APP_ID,
             node_id="test_node",
@@ -225,7 +243,12 @@ class TestWorkflowDraftVariableFields:
                 "length": 10,
                 "download_url": "http://example.com/signed-url",
             }
-            assert dump_response(WorkflowDraftVariableResponse, node_var) == expected_with_value
+            assert (
+                dump_response(
+                    WorkflowDraftVariableResponse, draft_variable_response_source(node_var, session=sqlite_session)
+                )
+                == expected_with_value
+            )
 
 
 class TestWorkflowDraftVariableList:
@@ -307,11 +330,42 @@ class TestWorkflowDraftVariableList:
             )
 
 
-def test_workflow_node_variables_fields() -> None:
+def test_value_response_rejects_bare_model() -> None:
+    """A bare model would expose the raw serialized JSON string as the value, so it must fail loudly."""
     conv_var = WorkflowDraftVariable.new_conversation_variable(
         app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1)
     )
-    resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[conv_var]))
+    with pytest.raises(TypeError, match="draft_variable_response_source"):
+        dump_response(WorkflowDraftVariableResponse, conv_var)
+
+
+def test_value_response_accepts_plain_mapping() -> None:
+    resp = WorkflowDraftVariableResponse.model_validate(
+        {
+            "id": "variable-id",
+            "type": "conversation",
+            "name": "conv_var",
+            "description": "",
+            "selector": ["conversation", "conv_var"],
+            "value_type": "number",
+            "edited": False,
+            "visible": True,
+            "is_truncated": False,
+            "value": 1,
+            "full_content": None,
+        }
+    )
+    assert resp.value == 1
+
+
+def test_workflow_node_variables_fields(sqlite_session: Session) -> None:
+    conv_var = WorkflowDraftVariable.new_conversation_variable(
+        app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1)
+    )
+    resp = dump_response(
+        WorkflowDraftVariableListResponse,
+        draft_variable_list_response_source(WorkflowDraftVariableList(variables=[conv_var]), session=sqlite_session),
+    )
     assert isinstance(resp, dict)
     assert len(resp["items"]) == 1
     item_dict = resp["items"][0]
@@ -319,7 +373,7 @@ def test_workflow_node_variables_fields() -> None:
     assert item_dict["value"] == 1
 
 
-def test_workflow_file_variable_with_signed_url() -> None:
+def test_workflow_file_variable_with_signed_url(sqlite_session: Session) -> None:
     """Test that File type variables include signed URLs in API responses."""
     from graphon.file import File, FileTransferMethod, FileType
 
@@ -345,7 +399,10 @@ def test_workflow_file_variable_with_signed_url() -> None:
     )
 
     # Marshal the variable using the API fields
-    resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[file_var]))
+    resp = dump_response(
+        WorkflowDraftVariableListResponse,
+        draft_variable_list_response_source(WorkflowDraftVariableList(variables=[file_var]), session=sqlite_session),
+    )
 
     # Verify the response structure
     assert isinstance(resp, dict)
@@ -375,7 +432,7 @@ def test_workflow_file_variable_with_signed_url() -> None:
     assert "sign=" in remote_url
 
 
-def test_workflow_file_variable_remote_url() -> None:
+def test_workflow_file_variable_remote_url(sqlite_session: Session) -> None:
     """Test that File type variables with REMOTE_URL transfer method return the remote URL."""
     from graphon.file import File, FileTransferMethod, FileType
 
@@ -401,7 +458,10 @@ def test_workflow_file_variable_remote_url() -> None:
     )
 
     # Marshal the variable using the API fields
-    resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[file_var]))
+    resp = dump_response(
+        WorkflowDraftVariableListResponse,
+        draft_variable_list_response_source(WorkflowDraftVariableList(variables=[file_var]), session=sqlite_session),
+    )
 
     # Verify the response structure
     assert isinstance(resp, dict)

api/tests/unit_tests/models/test_workflow.py

@@ -36,7 +36,7 @@ from models import Workflow as _workflow_feature_compatibility_Workflow
 from models.account import Account
 from models.base import Base, TypeBase
 from models.enums import CreatorUserRole, EndUserType, ExecutionOffLoadType, WorkflowRunTriggeredFrom
-from models.model import EndUser, UploadFile
+from models.model import App, AppMode, EndUser, UploadFile
 from models.tools import WorkflowToolProvider
 from models.workflow import (
     Workflow as _workflow_Workflow,
@@ -54,6 +54,7 @@ from models.workflow import (
     WorkflowRun,
     WorkflowType,
     WorkflowVersionCounter,
+    _resolve_workflow_app_tenant_id,
     is_system_variable_editable,
 )
 from models.workflow import WorkflowType as ModelWorkflowType
@@ -379,7 +380,7 @@ class TestWorkflowCases:
             assert is_system_variable_editable("invalid_or_new_system_variable") is False
 
     class TestWorkflowDraftVariableGetValue:
-        def test_get_value_by_case(self) -> None:
+        def test_get_value_by_case(self, sqlite_session: Session) -> None:
             @dataclasses.dataclass
             class TestCase:
                 name: str
@@ -449,9 +450,9 @@ class TestWorkflowCases:
                 fail_msg = f"test case {c.name} failed, index={idx}"
                 draft_var = WorkflowDraftVariable()
                 draft_var.set_value(c.value)
-                assert c.value == draft_var.get_value(), fail_msg
+                assert c.value == draft_var.get_value(session=sqlite_session), fail_msg
 
-        def test_file_variable_preserves_all_fields(self) -> None:
+        def test_file_variable_preserves_all_fields(self, sqlite_session: Session) -> None:
             """Test that File type variables preserve all fields during encoding/decoding."""
 
             # Create a File with specific field values
@@ -473,7 +474,7 @@ class TestWorkflowCases:
             draft_var.set_value(file_segment)
 
             # Retrieve the value and verify all fields are preserved
-            retrieved_segment = draft_var.get_value()
+            retrieved_segment = draft_var.get_value(session=sqlite_session)
             retrieved_file = retrieved_segment.value
 
             # Verify all important fields are preserved
@@ -490,7 +491,7 @@ class TestWorkflowCases:
             # Verify the segments have the same type and the important fields match
             assert file_segment.value_type == retrieved_segment.value_type
 
-        def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self) -> None:
+        def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self, sqlite_session: Session) -> None:
             persisted_file = File(
                 file_id="test_file_id",
                 file_type=FileType.DOCUMENT,
@@ -522,17 +523,58 @@ class TestWorkflowCases:
                 mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"),
                 mock.patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file,
             ):
-                retrieved_segment = draft_var.get_value()
+                retrieved_segment = draft_var.get_value(session=sqlite_session)
 
             assert retrieved_segment.value == rebuilt_file
             rebuild_file.assert_called_once()
             assert rebuild_file.call_args.kwargs["tenant_id"] == "tenant-1"
 
-        def test_get_and_set_value(self) -> None:
+        def test_array_file_variable_rebuilds_each_stored_payload_with_app_tenant(
+            self, sqlite_session: Session
+        ) -> None:
+            files = [
+                File(
+                    file_id=f"test_file_id_{index}",
+                    file_type=FileType.DOCUMENT,
+                    transfer_method=FileTransferMethod.LOCAL_FILE,
+                    reference=build_file_reference(record_id=f"upload-{index}", storage_key="legacy-storage-key"),
+                    filename=f"test_{index}.txt",
+                    extension=".txt",
+                    mime_type="text/plain",
+                    size=12,
+                )
+                for index in range(2)
+            ]
+            rebuilt_files = [
+                file.model_copy(update={"storage_key": "canonical-storage-key"}, deep=True) for file in files
+            ]
+            encoded = WorkflowDraftVariable(
+                app_id="app-1",
+            )
+            encoded.set_value(build_segment(files))
+            # A fresh instance holding only the persisted columns, so get_value decodes from scratch.
+            draft_var = WorkflowDraftVariable(
+                app_id="app-1",
+            )
+            draft_var.value = encoded.value
+            draft_var.value_type = encoded.value_type
+
+            with (
+                mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1") as resolve,
+                mock.patch("models.workflow.build_file_from_stored_mapping", side_effect=rebuilt_files) as rebuild_file,
+            ):
+                retrieved_segment = draft_var.get_value(session=sqlite_session)
+
+            assert retrieved_segment.value == rebuilt_files
+            resolve.assert_called_once_with(sqlite_session, "app-1")
+            assert rebuild_file.call_count == 2
+            assert all(call.kwargs["tenant_id"] == "tenant-1" for call in rebuild_file.call_args_list)
+
+        def test_get_and_set_value(self, sqlite_session: Session) -> None:
             draft_var = WorkflowDraftVariable()
             int_var = IntegerSegment(value=1)
             draft_var.set_value(int_var)
-            value = draft_var.get_value()
+            value = draft_var.get_value(session=sqlite_session)
             assert value == int_var
 
 
@@ -2177,3 +2219,22 @@ class TestWorkflowTypeBaseCases:
         relationship = inspect(WorkflowNodeExecutionModel).relationships["offload_data"]
         assert relationship.mapper.class_ is WorkflowNodeExecutionOffload
         assert relationship.back_populates == "execution"
+
+
+@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
+def test_resolve_workflow_app_tenant_id_reads_through_caller_session(sqlite_session: Session) -> None:
+    app = App(
+        id="app-1",
+        tenant_id="tenant-1",
+        name="Workflow",
+        mode=AppMode.WORKFLOW,
+        enable_site=True,
+        enable_api=True,
+    )
+    sqlite_session.add(app)
+    sqlite_session.flush()
+
+    assert _resolve_workflow_app_tenant_id(sqlite_session, "app-1") == "tenant-1"
+
+    with pytest.raises(ValueError, match="Unable to resolve tenant_id"):
+        _resolve_workflow_app_tenant_id(sqlite_session, "missing-app")

api/tests/unit_tests/services/test_workflow_draft_variable_service.py

@@ -27,7 +27,7 @@ from graphon.enums import BuiltinNodeTypes as _workflow_draft_variable_service_B
 from graphon.enums import WorkflowNodeExecutionStatus
 from graphon.file import File, FileTransferMethod, FileType
 from graphon.nodes import BuiltinNodeTypes as _workflow_draft_variable_service_sqlite_BuiltinNodeTypes
-from graphon.variables.segments import StringSegment
+from graphon.variables.segments import ArrayFileSegment, FileSegment, StringSegment
 from graphon.variables.types import SegmentType
 from graphon.variables.variables import StringVariable
 from libs import datetime_utils
@@ -216,6 +216,70 @@ class TestDraftVarLoaderSimpleCases:
             with pytest.raises(AssertionError):
                 draft_var_loader._load_offloaded_variable(draft_var)
 
+        @pytest.mark.parametrize(
+            "sqlite_session",
+            [(WorkflowDraftVariable, WorkflowDraftVariableFile, UploadFile)],
+            indirect=True,
+        )
+        def test_load_variables_collects_file_values_for_storage_key_loading(
+            self,
+            draft_var_loader: DraftVarLoader,
+            sqlite_session: Session,
+        ):
+            """File and array-file variables surface their File objects to the storage-key loader."""
+
+            def _make_file(index: int) -> File:
+                return File(
+                    file_id=f"file-{index}",
+                    file_type=FileType.DOCUMENT,
+                    transfer_method=FileTransferMethod.LOCAL_FILE,
+                    reference=build_file_reference(record_id=f"upload-{index}", storage_key="legacy-storage-key"),
+                    filename=f"test_{index}.txt",
+                    extension=".txt",
+                    mime_type="text/plain",
+                    size=12,
+                )
+
+            single_file = _make_file(0)
+            array_files = [_make_file(1), _make_file(2)]
+            file_var = WorkflowDraftVariable.new_node_variable(
+                app_id="test-app-id",
+                user_id="test-user-id",
+                node_id="node1",
+                name="file_var",
+                value=FileSegment(value=single_file),
+                node_execution_id="execution-node1",
+            )
+            array_var = WorkflowDraftVariable.new_node_variable(
+                app_id="test-app-id",
+                user_id="test-user-id",
+                node_id="node1",
+                name="array_var",
+                value=ArrayFileSegment(value=array_files),
+                node_execution_id="execution-node1",
+            )
+            sqlite_session.add_all([file_var, array_var])
+            sqlite_session.commit()
+
+            rebuilt = {
+                file.id: file.model_copy(update={"storage_key": "canonical"}, deep=True)
+                for file in [single_file, *array_files]
+            }
+
+            with (
+                patch("models.workflow._resolve_workflow_app_tenant_id", return_value="test-tenant-id"),
+                patch(
+                    "models.workflow.build_file_from_stored_mapping",
+                    side_effect=lambda **kwargs: rebuilt[kwargs["file_mapping"]["id"]],
+                ),
+                patch("services.workflow_draft_variable_service.StorageKeyLoader") as storage_key_loader_cls,
+            ):
+                result = draft_var_loader.load_variables([["node1", "file_var"], ["node1", "array_var"]])
+
+            loaded_files = storage_key_loader_cls.return_value.load_storage_keys.call_args.args[0]
+            assert {file.id for file in loaded_files} == {"file-0", "file-1", "file-2"}
+            assert {variable.name for variable in result} == {"file_var", "array_var"}
+
         def test_load_variables_empty_selectors_unit(self, draft_var_loader):
             """Test load_variables returns empty list for empty selectors."""
             result = draft_var_loader.load_variables([])
@@ -616,7 +680,7 @@ class TestWorkflowDraftVariableServiceCases:
                 draft_vars = saver._build_variables_from_start_mapping({"sys.files": [raw_file]})
 
             sys_var = draft_vars[0]
-            assert sys_var.get_value().value[0] == rebuilt_file
+            assert sys_var.get_value(session=sqlite_session).value[0] == rebuilt_file
             rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1")
 
         @pytest.fixture
@@ -1307,7 +1371,7 @@ class TestWorkflowDraftVariableServiceSqliteCases:
             assert node_var is not None
             assert node_var.id == self._node1_str_var_id
             assert node_var.name == "str_var"
-            assert node_var.get_value() == build_segment("str_value")
+            assert node_var.get_value(session=self._session) == build_segment("str_value")
 
         def test_get_system_variable(self) -> None:
             srv = self._get_test_srv()
@@ -1315,7 +1379,7 @@ class TestWorkflowDraftVariableServiceSqliteCases:
             assert sys_var is not None
             assert sys_var.id == self._sys_var_id
             assert sys_var.name == "sys_var"
-            assert sys_var.get_value() == build_segment("sys_value")
+            assert sys_var.get_value(session=self._session) == build_segment("sys_value")
 
         def test_get_conversation_variable(self) -> None:
             srv = self._get_test_srv()
@@ -1323,7 +1387,7 @@ class TestWorkflowDraftVariableServiceSqliteCases:
             assert conv_var is not None
             assert conv_var.id == self._conv_var_id
             assert conv_var.name == "conv_var"
-            assert conv_var.get_value() == build_segment("conv_value")
+            assert conv_var.get_value(session=self._session) == build_segment("conv_value")
 
         def test_delete_node_variables(self) -> None:
             srv = self._get_test_srv()

api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py

@@ -2,11 +2,9 @@ from __future__ import annotations
 
 import json
 from datetime import UTC, datetime
-from types import SimpleNamespace
 from typing import cast
 from unittest.mock import MagicMock
 
-import pytest
 from sqlalchemy.orm import Session, sessionmaker
 
 from domain.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt
@@ -18,13 +16,6 @@ from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTri
 from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository
 from services.workflow_node_execution_trace_service import assemble_workflow_node_execution_traces
 
-pytestmark = pytest.mark.parametrize("sqlite_session", [(Account,)], indirect=True)
-
-
-@pytest.fixture(autouse=True)
-def _bind_model_database(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None:
-    monkeypatch.setattr("models.workflow.db", SimpleNamespace(session=sqlite_session))
-
 
 def _retry_attempt(retry_index: int, **overrides: object) -> dict[str, object]:
     values = {