2026-08-20 11:36:27 +08:00
5 changed files with 267 additions and 17 deletions
--- a/easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/easyagentsflow/code/PythonRuntimeEngine.java
+++ b/easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/easyagentsflow/code/PythonRuntimeEngine.java
@@ -14,10 +14,19 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;

+/**
+ * 基于独立 Python 子进程的代码节点执行器。
+ *
+ * <p>优先保留历史 {@code _result} 输出约定；当脚本未写入
+ * {@code _result} 且声明了 {@code main} 函数时，自动传入节点参数并使用
+ * {@code main} 返回的字典作为节点输出。</p>
+ */
 public class PythonRuntimeEngine implements CodeRuntimeEngine {

    private static final String PYTHON_HELPER_SCRIPT = String.join("\n",
@@ -27,6 +36,40 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
            "import sys",
            "import traceback",
            "",
+            "class _BoundedTextWriter(io.TextIOBase):",
+            "    def __init__(self, max_bytes):",
+            "        super().__init__()",
+            "        self._max_bytes = max(0, max_bytes)",
+            "        self._written_bytes = 0",
+            "        self._buffer = io.StringIO()",
+            "        self.truncated = False",
+            "",
+            "    def write(self, text):",
+            "        if not isinstance(text, str):",
+            "            raise TypeError('write() argument must be str')",
+            "        if not text:",
+            "            return 0",
+            "        remaining = self._max_bytes - self._written_bytes",
+            "        if remaining <= 0:",
+            "            self.truncated = True",
+            "            return len(text)",
+            "        candidate = text[:remaining]",
+            "        encoded = candidate.encode('utf-8')",
+            "        if len(encoded) > remaining:",
+            "            candidate = encoded[:remaining].decode('utf-8', errors='ignore')",
+            "            encoded = candidate.encode('utf-8')",
+            "        self._buffer.write(candidate)",
+            "        self._written_bytes += len(encoded)",
+            "        if len(candidate) < len(text):",
+            "            self.truncated = True",
+            "        return len(text)",
+            "",
+            "    def getvalue(self):",
+            "        return self._buffer.getvalue()",
+            "",
+            "    def writable(self):",
+            "        return True",
+            "",
            "def _read_payload():",
            "    raw = sys.stdin.read()",
            "    if not raw:",
@@ -39,22 +82,39 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
            "    context = payload.get('context')",
            "    if not isinstance(context, dict):",
            "        context = {}",
+            "    parameter_names = payload.get('parameterNames')",
+            "    if not isinstance(parameter_names, list):",
+            "        parameter_names = []",
+            "    max_output_bytes = payload.get('maxOutputBytes')",
+            "    if not isinstance(max_output_bytes, int):",
+            "        max_output_bytes = 0",
            "",
            "    local_ctx = dict(context)",
            "    local_ctx['_result'] = {}",
+            "    local_ctx['__builtins__'] = __builtins__",
            "",
-            "    captured_stdout = io.StringIO()",
-            "    captured_stderr = io.StringIO()",
+            "    captured_stdout = _BoundedTextWriter(max_output_bytes)",
+            "    captured_stderr = _BoundedTextWriter(max_output_bytes)",
+            "    main_invoked = False",
            "",
            "    try:",
            "        with contextlib.redirect_stdout(captured_stdout), contextlib.redirect_stderr(captured_stderr):",
-            "            exec(code, {'__builtins__': __builtins__}, local_ctx)",
+            "            exec(code, local_ctx, local_ctx)",
+            "            result = local_ctx.get('_result')",
+            "            if isinstance(result, dict) and not result:",
+            "                main_func = local_ctx.get('main')",
+            "                if callable(main_func):",
+            "                    main_invoked = True",
+            "                    main_input = {name: context.get(name) for name in parameter_names if isinstance(name, str)}",
+            "                    result = main_func(main_input)",
            "    except Exception:",
            "        sys.stderr.write(traceback.format_exc())",
            "        sys.exit(2)",
            "",
-            "    result = local_ctx.get('_result')",
            "    if not isinstance(result, dict):",
+            "        if main_invoked:",
+            "            sys.stderr.write(\"Python main 函数必须返回 dict，例如：return {'result': value}\\n\")",
+            "        else:",
            "            sys.stderr.write('Python 脚本执行结束后，_result 必须是 dict\\n')",
            "        sys.exit(3)",
            "",
@@ -130,8 +190,37 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
        }
    }

+    /**
+     * 执行 Python 脚本，兼容直接调用旧接口的场景。
+     *
+     * @param code Python 脚本
+     * @param node 当前代码节点
+     * @param chain 当前工作流
+     * @return 脚本输出结果
+     */
    @Override
    public Map<String, Object> execute(String code, CodeNode node, Chain chain) {
+        Map<String, Object> parameterValues =
+                chain.getExecutionState().resolveParameters(node);
+        return execute(code, node, chain, parameterValues);
+    }
+
+    /**
+     * 使用调用方已解析的参数执行 Python 脚本。
+     *
+     * <p>该方法作为普通重载保留，可同时兼容仅声明三参数接口的旧版
+     * easy-agents；新版接口会按同签名方法直接分派，避免重复解析参数。</p>
+     *
+     * @param code Python 脚本
+     * @param node 当前代码节点
+     * @param chain 当前工作流
+     * @param parameterValues 已解析的节点参数
+     * @return 脚本输出结果
+     */
+    public Map<String, Object> execute(String code,
+                                       CodeNode node,
+                                       Chain chain,
+                                       Map<String, Object> parameterValues) {
        if (StringUtil.noText(command)) {
            throw new RuntimeException("python 执行命令为空，请检查配置 node.code-engine.python.command");
        }
@@ -158,10 +247,16 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
            stdoutThread.start();
            stderrThread.start();

-            Map<String, Object> context = buildContext(chain, node);
+            Map<String, Object> context = buildContext(chain, parameterValues);
            Map<String, Object> payload = new HashMap<>();
            payload.put("code", code);
            payload.put("context", context);
+            payload.put(
+                    "parameterNames",
+                    parameterValues == null
+                            ? Collections.emptyList()
+                            : new ArrayList<>(parameterValues.keySet()));
+            payload.put("maxOutputBytes", maxOutputBytes);

            String payloadText = JSON.toJSONString(payload);
            try (OutputStream stdin = process.getOutputStream()) {
@@ -217,7 +312,15 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
        }
    }

-    private Map<String, Object> buildContext(Chain chain, CodeNode node) {
+    /**
+     * 构建传入 Python 子进程的上下文。
+     *
+     * @param chain 当前工作流
+     * @param parameterValues 已解析的节点参数
+     * @return Python 脚本上下文
+     */
+    private Map<String, Object> buildContext(
+            Chain chain, Map<String, Object> parameterValues) {
        Map<String, Object> context = new HashMap<>();
        ChainState chainState =
                chain.getExecutionState();
@@ -230,8 +333,6 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
            }
        });

-        Map<String, Object> parameterValues =
-                chainState.resolveParameters(node);
        if (parameterValues != null && !parameterValues.isEmpty()) {
            context.putAll(parameterValues);
        }
--- a/easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/easyagentsflow/code/PythonRuntimeEngineTest.java
+++ b/easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/easyagentsflow/code/PythonRuntimeEngineTest.java
@@ -31,6 +31,106 @@ public class PythonRuntimeEngineTest {
        Assert.assertEquals("dict", result.get("env_type"));
    }

+    @Test
+    public void testExecuteMainSuccess() {
+        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
+
+        Chain chain = createChain();
+        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
+        chain.getState().getMemory().put("ambient", "memory-value");
+        PythonRuntimeEngine engine = new PythonRuntimeEngine(
+                PYTHON_COMMAND,
+                3000L,
+                65536,
+                System.getProperty("java.io.tmpdir"));
+
+        Map<String, Object> result = engine.execute(
+                String.join("\n",
+                        "import math",
+                        "def main(inputs):",
+                        "    return {",
+                        "        'answer': inputs['input'],",
+                        "        'score': 95,",
+                        "        'sqrt': math.sqrt(16),",
+                        "        'has_ambient': 'ambient' in inputs",
+                        "    }"),
+                node,
+                chain,
+                Map.of("input", "hello"));
+
+        Assert.assertEquals("hello", result.get("answer"));
+        Assert.assertEquals(95, ((Number) result.get("score")).intValue());
+        Assert.assertEquals(4D, ((Number) result.get("sqrt")).doubleValue(), 0D);
+        Assert.assertEquals(Boolean.FALSE, result.get("has_ambient"));
+    }
+
+    @Test
+    public void testLegacyResultTakesPrecedenceOverMain() {
+        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
+
+        Chain chain = createChain();
+        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
+        PythonRuntimeEngine engine = new PythonRuntimeEngine(
+                PYTHON_COMMAND,
+                3000L,
+                65536,
+                System.getProperty("java.io.tmpdir"));
+
+        Map<String, Object> result = engine.execute(
+                String.join("\n",
+                        "_result['source'] = 'legacy'",
+                        "def main(inputs):",
+                        "    return {'source': 'main'}"),
+                node,
+                chain,
+                Map.of("input", "hello"));
+
+        Assert.assertEquals("legacy", result.get("source"));
+    }
+
+    @Test
+    public void testMainMustReturnDict() {
+        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
+
+        Chain chain = createChain();
+        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
+        PythonRuntimeEngine engine = new PythonRuntimeEngine(
+                PYTHON_COMMAND,
+                3000L,
+                65536,
+                System.getProperty("java.io.tmpdir"));
+
+        assertExecuteFail(
+                engine,
+                node,
+                chain,
+                "def main(inputs):\n    return 'invalid'",
+                "Python main 函数必须返回 dict");
+    }
+
+    @Test
+    public void testCapturedOutputIsBounded() {
+        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
+
+        Chain chain = createChain();
+        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
+        PythonRuntimeEngine engine = new PythonRuntimeEngine(
+                PYTHON_COMMAND,
+                3000L,
+                64,
+                System.getProperty("java.io.tmpdir"));
+
+        Map<String, Object> result = engine.execute(
+                String.join("\n",
+                        "import sys",
+                        "print('a' * 2048)",
+                        "_result['truncated'] = sys.stdout.truncated"),
+                node,
+                chain);
+
+        Assert.assertEquals(Boolean.TRUE, result.get("truncated"));
+    }
+
    @Test
    public void testSyntaxError() {
        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
--- a/easyflow-ui-admin/packages/tinyflow-ui/src/components/nodes/CodeNode.svelte
+++ b/easyflow-ui-admin/packages/tinyflow-ui/src/components/nodes/CodeNode.svelte
@@ -22,8 +22,8 @@
    let currentNode = useNodesData(currentNodeId);
    const { addParameter } = useAddParameter();
    const { updateNodeData } = useSvelteFlow();
-    const codeNodeHelp = `JavaScript 如何返回结果
- 定义 main 函数并返回对象，系统会自动执行。
+    const codeNodeHelp = `代码如何返回结果
+- JavaScript 定义 main 函数并返回对象；Python 定义 main 函数并返回 dict，系统会自动执行。
 - main 接收由输入参数组成的对象。

 输出参数如何配置
@@ -31,11 +31,15 @@
 - 字段名与 main 返回对象的 key 保持一致。
 - 下游节点可引用：代码节点ID.输出参数名。

-示例
- 代码里写：
+JavaScript 示例
  function main({ input }) {
    return { answer: input, score: 95 };
  }
+
+Python 示例
+  def main(inputs):
+      return {'answer': inputs.get('input'), 'score': 95}
+
 - 输出参数配置：answer(String)、score(Number)
 - 结束节点输出参数可引用：代码节点ID.answer、代码节点ID.score

@@ -53,6 +57,13 @@
        const firstAvailable = engines.find((item) => item.selectable !== false);
        return firstAvailable?.value || 'js';
    });
+    const codePlaceholder = $derived.by(() => {
+        const engine = String(data.engine || defaultEngine).trim().toLowerCase();
+        if (engine === 'python' || engine === 'py') {
+            return "请输入代码，例如：def main(inputs):\n    return {'output': inputs.get('input')}";
+        }
+        return '请输入代码，例如：function main({ input }) { return { output: input }; }';
+    });

    onMount(async () => {
        const codeEngines = await options.provider?.codeEngine?.();
@@ -118,7 +129,7 @@
        <CodeScriptEditor
                  mode="textarea"
                  rows={10}
-                  placeholder={'请输入代码，例如：function main({ input }) { return { output: input }; }'}
+                  placeholder={codePlaceholder}
                  style="width: 100%"
                  engine={(data.engine as string) || (defaultEngine as string)}
                  parameters={editorParameters}
--- a/easyflow-ui-admin/packages/tinyflow-ui/src/components/utils/codeCompletion.test.ts
+++ b/easyflow-ui-admin/packages/tinyflow-ui/src/components/utils/codeCompletion.test.ts
@@ -71,6 +71,36 @@ describe('codeCompletion utils', () => {
    );
  });

+  it('should prefer main return snippets for python output', () => {
+    const completions = createBusinessCompletions('python', []);
+    const mainCompletion = completions.find((item) => item.label === 'main');
+    const resultCompletion = completions.find(
+      (item) => item.label === 'result-object',
+    );
+
+    expect(mainCompletion).toMatchObject({
+      type: 'snippet',
+      detail: '代码节点入口',
+      boost: 1000,
+    });
+    expect(String(mainCompletion?.apply)).toContain('def main(inputs):');
+    expect(String(mainCompletion?.apply)).toContain('return {');
+    expect(String(resultCompletion?.apply)).toContain('def main(inputs):');
+    expect(String(resultCompletion?.apply)).toContain(
+      "return {'message': 'ok', 'data': inputs.get('data')}",
+    );
+  });
+
+  it('should provide standalone python snippets', () => {
+    const completions = createBusinessCompletions('python', []);
+    const snippets = completions.filter((item) => item.type === 'snippet');
+
+    expect(snippets).toHaveLength(4);
+    for (const snippet of snippets) {
+      expect(String(snippet.apply)).toMatch(/^def main\(inputs\):\n    /);
+    }
+  });
+
  it('should detect blocked nodes for comment/string/property contexts', () => {
    expect(shouldSkipBusinessCompletion('Comment')).toBe(true);
    expect(shouldSkipBusinessCompletion('LineComment')).toBe(true);
--- a/easyflow-ui-admin/packages/tinyflow-ui/src/components/utils/codeCompletion.ts
+++ b/easyflow-ui-admin/packages/tinyflow-ui/src/components/utils/codeCompletion.ts
@@ -60,21 +60,29 @@ const PYTHON_SNIPPETS: Array<{
  insert: string;
  detail: string;
 }> = [
+  {
+    label: 'main',
+    detail: '代码节点入口',
+    insert:
+      "def main(inputs):\n    return {\n        'output': inputs.get('input'),\n    }",
+  },
  {
    label: 'if-else',
    detail: '条件分支',
    insert:
-      "if condition:\n    _result['value'] = value\nelse:\n    _result['value'] = None",
+      "def main(inputs):\n    condition = inputs.get('condition')\n    value = inputs.get('value')\n    if condition:\n        return {'value': value}\n    return {'value': None}",
  },
  {
    label: 'for-loop',
    detail: '遍历数组',
-    insert: "for item in items:\n    # TODO\n    pass\n_result['done'] = True",
+    insert:
+      "def main(inputs):\n    result = []\n    for item in inputs.get('items', []):\n        result.append(item)\n    return {'result': result}",
  },
  {
    label: 'result-object',
    detail: '返回对象',
-    insert: "_result['message'] = 'ok'\n_result['data'] = data",
+    insert:
+      "def main(inputs):\n    return {'message': 'ok', 'data': inputs.get('data')}",
  },
 ];