From 09605dec89683b3e19743be4c82aa7fc76d0bf31 Mon Sep 17 00:00:00 2001
From: QuantumGhost <obelisk.reg+git@gmail.com>
Date: Wed, 25 Jun 2025 11:26:23 +0800
Subject: [PATCH] fix(api): fix the issue that nested objects are not added to
 variable pool properly

---
 .../workflow/graph_engine/graph_engine.py     | 17 +++++------
 api/core/workflow/utils/variable_utils.py     | 28 +++++++++++++++++++
 api/core/workflow/variable_loader.py          |  7 ++++-
 .../workflow_draft_variable_service.py        |  3 +-
 4 files changed, 43 insertions(+), 12 deletions(-)
 create mode 100644 api/core/workflow/utils/variable_utils.py

diff --git a/api/core/workflow/graph_engine/graph_engine.py b/api/core/workflow/graph_engine/graph_engine.py
index 2809afad02..61a7a26652 100644
--- a/api/core/workflow/graph_engine/graph_engine.py
+++ b/api/core/workflow/graph_engine/graph_engine.py
@@ -53,6 +53,7 @@ from core.workflow.nodes.end.end_stream_processor import EndStreamProcessor
 from core.workflow.nodes.enums import ErrorStrategy, FailBranchSourceHandle
 from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent
 from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
+from core.workflow.utils import variable_utils
 from libs.flask_utils import preserve_flask_contexts
 from models.enums import UserFrom
 from models.workflow import WorkflowType
@@ -856,16 +857,12 @@ class GraphEngine:
         :param variable_value: variable value
         :return:
         """
-        self.graph_runtime_state.variable_pool.add([node_id] + variable_key_list, variable_value)
-
-        # if variable_value is a dict, then recursively append variables
-        if isinstance(variable_value, dict):
-            for key, value in variable_value.items():
-                # construct new key list
-                new_key_list = variable_key_list + [key]
-                self._append_variables_recursively(
-                    node_id=node_id, variable_key_list=new_key_list, variable_value=value
-                )
+        variable_utils.append_variables_recursively(
+            self.graph_runtime_state.variable_pool,
+            node_id,
+            variable_key_list,
+            variable_value,
+        )
 
     def _is_timed_out(self, start_at: float, max_execution_time: int) -> bool:
         """
diff --git a/api/core/workflow/utils/variable_utils.py b/api/core/workflow/utils/variable_utils.py
new file mode 100644
index 0000000000..e68d990b60
--- /dev/null
+++ b/api/core/workflow/utils/variable_utils.py
@@ -0,0 +1,28 @@
+from core.variables.segments import ObjectSegment, Segment
+from core.workflow.entities.variable_pool import VariablePool, VariableValue
+
+
+def append_variables_recursively(
+    pool: VariablePool, node_id: str, variable_key_list: list[str], variable_value: VariableValue | Segment
+):
+    """
+    Append variables recursively
+    :param node_id: node id
+    :param variable_key_list: variable key list
+    :param variable_value: variable value
+    :return:
+    """
+    pool.add([node_id] + variable_key_list, variable_value)
+
+    # if variable_value is a dict, then recursively append variables
+    if isinstance(variable_value, ObjectSegment):
+        variable_dict = variable_value.value
+    elif isinstance(variable_value, dict):
+        variable_dict = variable_value
+    else:
+        return
+
+    for key, value in variable_dict.items():
+        # construct new key list
+        new_key_list = variable_key_list + [key]
+        append_variables_recursively(pool, node_id=node_id, variable_key_list=new_key_list, variable_value=value)
diff --git a/api/core/workflow/variable_loader.py b/api/core/workflow/variable_loader.py
index 4842ee00a5..1e13871d0a 100644
--- a/api/core/workflow/variable_loader.py
+++ b/api/core/workflow/variable_loader.py
@@ -3,7 +3,9 @@ from collections.abc import Mapping, Sequence
 from typing import Any, Protocol
 
 from core.variables import Variable
+from core.variables.consts import MIN_SELECTORS_LENGTH
 from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.utils import variable_utils
 
 
 class VariableLoader(Protocol):
@@ -76,4 +78,7 @@ def load_into_variable_pool(
             variables_to_load.append(list(selector))
     loaded = variable_loader.load_variables(variables_to_load)
     for var in loaded:
-        variable_pool.add(var.selector, var)
+        assert len(var.selector) >= MIN_SELECTORS_LENGTH, f"Invalid variable {var}"
+        variable_utils.append_variables_recursively(
+            variable_pool, node_id=var.selector[0], variable_key_list=list(var.selector[1:]), variable_value=var
+        )
diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py
index cd30440b4f..164693c2e1 100644
--- a/api/services/workflow_draft_variable_service.py
+++ b/api/services/workflow_draft_variable_service.py
@@ -129,7 +129,8 @@ class WorkflowDraftVariableService:
     ) -> list[WorkflowDraftVariable]:
         ors = []
         for selector in selectors:
-            node_id, name = selector
+            assert len(selector) >= MIN_SELECTORS_LENGTH, f"Invalid selector to get: {selector}"
+            node_id, name = selector[:2]
             ors.append(and_(WorkflowDraftVariable.node_id == node_id, WorkflowDraftVariable.name == name))
 
         # NOTE(QuantumGhost): Although the number of `or` expressions may be large, as long as