feat(code_node): add more check (#11949)

Signed-off-by: -LAN- <[email protected]>

api/core/helper/code_executor/code_executor.py

@@ -118,7 +118,7 @@ def execute_code(cls, language: CodeLanguage, preload: str, code: str) -> str:
         return response.data.stdout or ""
 
     @classmethod
-    def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]) -> dict:
+    def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]):
         """
         Execute code
         :param language: code language

api/core/helper/code_executor/template_transformer.py

@@ -25,23 +25,28 @@ def transform_caller(cls, code: str, inputs: Mapping[str, Any]) -> tuple[str, st
         return runner_script, preload_script
 
     @classmethod
-    def extract_result_str_from_response(cls, response: str) -> str:
+    def extract_result_str_from_response(cls, response: str):
         result = re.search(rf"{cls._result_tag}(.*){cls._result_tag}", response, re.DOTALL)
         if not result:
             raise ValueError("Failed to parse result")
         result = result.group(1)
         return result
 
     @classmethod
-    def transform_response(cls, response: str):
+    def transform_response(cls, response: str) -> Mapping[str, Any]:
         """
         Transform response to dict
         :param response: response
         :return:
         """
-        result = json.loads(cls.extract_result_str_from_response(response))
+        try:
+            result = json.loads(cls.extract_result_str_from_response(response))
+        except json.JSONDecodeError:
+            raise ValueError("failed to parse response")
         if not isinstance(result, dict):
-            raise ValueError("Result must be a dict")
+            raise ValueError("result must be a dict")
+        if not all(isinstance(k, str) for k in result):
+            raise ValueError("result keys must be strings")
         return result
 
     @classmethod

api/core/workflow/nodes/code/code_node.py

@@ -59,7 +59,7 @@ def _run(self) -> NodeRunResult:
             )
 
             # Transform result
-            result = self._transform_result(result, self.node_data.outputs)
+            result = self._transform_result(result=result, output_schema=self.node_data.outputs)
         except (CodeExecutionError, CodeNodeError) as e:
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error=str(e), error_type=type(e).__name__
@@ -116,14 +116,12 @@ def _check_number(self, value: int | float | None, variable: str) -> int | float
         return value
 
     def _transform_result(
-        self, result: dict, output_schema: Optional[dict[str, CodeNodeData.Output]], prefix: str = "", depth: int = 1
-    ) -> dict:
-        """
-        Transform result
-        :param result: result
-        :param output_schema: output schema
-        :return:
-        """
+        self,
+        result: Mapping[str, Any],
+        output_schema: Optional[dict[str, CodeNodeData.Output]],
+        prefix: str = "",
+        depth: int = 1,
+    ):
         if depth > dify_config.CODE_MAX_DEPTH:
             raise DepthLimitError(f"Depth limit ${dify_config.CODE_MAX_DEPTH} reached, object too deep.")