- feat(graph-engine): enforce variable update value types with SegmentType validation and casting

- test(graph-engine): update VariableUpdate usages to include value_type in command tests

Author: laipz8200

api/core/workflow/graph_engine/entities/commands.py

@@ -9,10 +9,10 @@
 from enum import StrEnum, auto
 from typing import Any, TypeAlias
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 from core.file import File
-from core.variables import Segment, Variable
+from core.variables import Segment, SegmentType, Variable
 
 
 class CommandType(StrEnum):
@@ -51,8 +51,31 @@ class VariableUpdate(BaseModel):
     """Represents a single variable update instruction."""
 
     selector: tuple[str, str] = Field(description="Variable selector (node_id, variable_name)")
+    value_type: SegmentType = Field(description="Variable value type")
     value: VariableUpdateValue = Field(description="New variable value")
 
+    @model_validator(mode="after")
+    def _validate_value_type(self) -> "VariableUpdate":
+        value_type = self.value_type
+        value = self.value
+
+        if isinstance(value, Variable | Segment):
+            if value.value_type != value_type:
+                raise ValueError(f"value type mismatch: expected {value_type}, got {value.value_type}")
+            return self
+
+        if isinstance(value, File):
+            if value_type != SegmentType.FILE:
+                raise ValueError(f"value type mismatch: expected {value_type}, got {SegmentType.FILE}")
+            return self
+
+        casted_value = SegmentType.cast_value(value, value_type)
+        if not value_type.is_valid(casted_value):
+            raise ValueError(f"value type mismatch: expected {value_type}, got {type(value).__name__}")
+
+        self.value = casted_value
+        return self
+
 
 class UpdateVariablesCommand(GraphEngineCommand):
     """Command to update a group of variables in the variable pool."""

api/tests/unit_tests/core/workflow/graph_engine/command_channels/test_redis_channel.py

@@ -3,6 +3,7 @@
 import json
 from unittest.mock import MagicMock
 
+from core.variables import SegmentType
 from core.workflow.graph_engine.command_channels.redis_channel import RedisChannel
 from core.workflow.graph_engine.entities.commands import (
     AbortCommand,
@@ -169,8 +170,8 @@ def test_fetch_commands_with_update_variables_command(self):
 
         update_command = UpdateVariablesCommand(
             updates=[
-                VariableUpdate(selector=["node1", "foo"], value="bar"),
-                VariableUpdate(selector=["node2", "baz"], value=123),
+                VariableUpdate(selector=["node1", "foo"], value_type=SegmentType.STRING, value="bar"),
+                VariableUpdate(selector=["node2", "baz"], value_type=SegmentType.INTEGER, value=123),
             ]
         )
         command_json = json.dumps(update_command.model_dump())

api/tests/unit_tests/core/workflow/graph_engine/test_command_system.py

@@ -4,6 +4,7 @@
 from unittest.mock import MagicMock
 
 from core.app.entities.app_invoke_entities import InvokeFrom
+from core.variables import SegmentType
 from core.workflow.entities.graph_init_params import GraphInitParams
 from core.workflow.entities.pause_reason import SchedulingPause
 from core.workflow.graph import Graph
@@ -231,8 +232,8 @@ def test_update_variables_command_updates_pool():
 
     update_command = UpdateVariablesCommand(
         updates=[
-            VariableUpdate(selector=["node1", "foo"], value="new value"),
-            VariableUpdate(selector=["node2", "bar"], value=123),
+            VariableUpdate(selector=["node1", "foo"], value_type=SegmentType.STRING, value="new value"),
+            VariableUpdate(selector=["node2", "bar"], value_type=SegmentType.INTEGER, value=123),
         ]
     )
     command_channel.send_command(update_command)