From 7e402abefd23692410558738f87b99dfbc7286ef Mon Sep 17 00:00:00 2001
From: QuantumGhost <obelisk.reg+git@gmail.com>
Date: Thu, 5 Jun 2025 20:54:30 +0800
Subject: [PATCH] fix(api): `set_value` should also update the cache

add relevant tests for `get_value` and `set_value` methods
---
 api/models/workflow.py                       |  9 ++++
 api/tests/unit_tests/models/test_workflow.py | 53 +++++++++++++++++++-
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/api/models/workflow.py b/api/models/workflow.py
index 28b661670e..9cb6968914 100644
--- a/api/models/workflow.py
+++ b/api/models/workflow.py
@@ -1027,6 +1027,15 @@ class WorkflowDraftVariable(Base):
         self._set_selector([self.node_id, name])
 
     def set_value(self, value: Segment):
+        """Updates the `value` and corresponding `value_type` fields in the database model.
+
+        This method also stores the provided Segment object in the deserialized cache
+        without creating a copy, allowing for efficient value access.
+
+        Args:
+            value: The Segment object to store as the variable's value.
+        """
+        self.__value = value
         self.value = json.dumps(value, cls=variable_utils.SegmentJSONEncoder)
         self.value_type = value.value_type
 
diff --git a/api/tests/unit_tests/models/test_workflow.py b/api/tests/unit_tests/models/test_workflow.py
index 04ebe7badc..69163d48bd 100644
--- a/api/tests/unit_tests/models/test_workflow.py
+++ b/api/tests/unit_tests/models/test_workflow.py
@@ -7,7 +7,7 @@ from constants import HIDDEN_VALUE
 from core.file.enums import FileTransferMethod, FileType
 from core.file.models import File
 from core.variables import FloatVariable, IntegerVariable, SecretVariable, StringVariable
-from core.variables.segments import Segment
+from core.variables.segments import IntegerSegment, Segment
 from factories.variable_factory import build_segment
 from models.workflow import Workflow, WorkflowDraftVariable, WorkflowNodeExecutionModel, is_system_variable_editable
 
@@ -262,4 +262,53 @@ class TestWorkflowDraftVariableGetValue:
             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(tenant_id=tenant_id), fail_msg
+            assert c.value == draft_var.get_value(), fail_msg
+
+    def test_file_variable_preserves_all_fields(self):
+        """Test that File type variables preserve all fields during encoding/decoding."""
+        tenant_id = "test_tenant_id"
+
+        # Create a File with specific field values
+        test_file = File(
+            id="test_file_id",
+            tenant_id=tenant_id,
+            type=FileType.IMAGE,
+            transfer_method=FileTransferMethod.REMOTE_URL,
+            remote_url="https://example.com/test.jpg",
+            filename="test.jpg",
+            extension=".jpg",
+            mime_type="image/jpeg",
+            size=12345,  # Specific size to test preservation
+            storage_key="test_storage_key",
+        )
+
+        # Create a FileSegment and WorkflowDraftVariable
+        file_segment = build_segment(test_file)
+        draft_var = WorkflowDraftVariable()
+        draft_var.set_value(file_segment)
+
+        # Retrieve the value and verify all fields are preserved
+        retrieved_segment = draft_var.get_value()
+        retrieved_file = retrieved_segment.value
+
+        # Verify all important fields are preserved
+        assert retrieved_file.id == test_file.id
+        assert retrieved_file.tenant_id == test_file.tenant_id
+        assert retrieved_file.type == test_file.type
+        assert retrieved_file.transfer_method == test_file.transfer_method
+        assert retrieved_file.remote_url == test_file.remote_url
+        assert retrieved_file.filename == test_file.filename
+        assert retrieved_file.extension == test_file.extension
+        assert retrieved_file.mime_type == test_file.mime_type
+        assert retrieved_file.size == test_file.size  # This was the main issue being fixed
+        # Note: storage_key is not serialized in model_dump() so it won't be preserved
+
+        # Verify the segments have the same type and the important fields match
+        assert file_segment.value_type == retrieved_segment.value_type
+
+    def test_get_and_set_value(self):
+        draft_var = WorkflowDraftVariable()
+        int_var = IntegerSegment(value=1)
+        draft_var.set_value(int_var)
+        value = draft_var.get_value()
+        assert value == int_var