zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-01 05:23:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refa: refine code_exec component (#13925)

Browse Source 
### What problem does this PR solve?

Refine code_exec component.

### Type of change

- [x] Refactoring
This commit is contained in:
Yongteng Lei
2026-04-07 11:48:29 +08:00
committed by GitHub
parent c4b0aaa874
commit 112007243d
30 changed files with 1767 additions and 427 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
agent/sandbox/executor_manager/models/schemas.py
View File

@@ -14,7 +14,7 @@
# limitations under the License.
#
import base64
from typing import Optional
from typing import Any, Optional
from pydantic import BaseModel, Field, field_validator
@@ -28,6 +28,12 @@ class ArtifactItem(BaseModel):
content_b64: str
class ExecutionStructuredResult(BaseModel):
present: bool
value: Any = None
type: str = "json"
class CodeExecutionResult(BaseModel):
status: ResultStatus
stdout: str
@@ -47,6 +53,9 @@ class CodeExecutionResult(BaseModel):
# File artifacts produced by code execution (images, PDFs, CSVs, etc.)
artifacts: list[ArtifactItem] = []
# Structured return value produced by main()
result: Optional[ExecutionStructuredResult] = None
class CodeExecutionRequest(BaseModel):
code_b64: str = Field(..., description="Base64 encoded code string")

105
agent/sandbox/executor_manager/services/execution.py
View File

@@ -19,14 +19,42 @@ import json
import os
import time
import uuid
from core.config import TIMEOUT
from core.container import allocate_container_blocking, release_container
from core.logger import logger
from models.enums import ResourceLimitType, ResultStatus, RuntimeErrorType, SupportLanguage, UnauthorizedAccessType
from models.schemas import ArtifactItem, CodeExecutionRequest, CodeExecutionResult
from models.schemas import ArtifactItem, CodeExecutionRequest, CodeExecutionResult, ExecutionStructuredResult
from utils.common import async_run_command
RESULT_MARKER_PREFIX = "__RAGFLOW_RESULT__:"
def _extract_result_envelope(stdout: str) -> tuple[str, ExecutionStructuredResult | None]:
if not stdout:
return "", None
cleaned_lines: list[str] = []
envelope: ExecutionStructuredResult | None = None
for line in str(stdout).splitlines():
if line.startswith(RESULT_MARKER_PREFIX):
payload_b64 = line[len(RESULT_MARKER_PREFIX) :].strip()
if not payload_b64:
continue
try:
payload = base64.b64decode(payload_b64).decode("utf-8")
envelope = ExecutionStructuredResult.model_validate_json(payload)
except Exception as exc:
logger.warning(f"Failed to decode structured result marker: {exc}")
cleaned_lines.append(line)
continue
cleaned_lines.append(line)
cleaned_stdout = "\n".join(cleaned_lines)
if stdout.endswith("\n") and cleaned_stdout and not cleaned_stdout.endswith("\n"):
cleaned_stdout += "\n"
return cleaned_stdout, envelope
async def execute_code(req: CodeExecutionRequest):
"""Fully asynchronous execution logic"""
@@ -48,15 +76,14 @@ async def execute_code(req: CodeExecutionRequest):
try:
if language == SupportLanguage.PYTHON:
code_name = "main.py"
# code
code_path = os.path.join(workdir, code_name)
with open(code_path, "wb") as f:
f.write(base64.b64decode(req.code_b64))
# runner
runner_name = "runner.py"
runner_path = os.path.join(workdir, runner_name)
with open(runner_path, "w") as f:
f.write("""import json
f.write(f"""import base64
import json
import os
import sys
@@ -65,33 +92,64 @@ os.makedirs(os.path.join(os.getcwd(), "artifacts"), exist_ok=True)
sys.path.insert(0, os.path.dirname(__file__))
from main import main
RESULT_MARKER_PREFIX = {RESULT_MARKER_PREFIX!r}
def emit_result(value):
payload = json.dumps(
{{
"present": True,
"value": value,
"type": "json",
}},
ensure_ascii=False,
separators=(",", ":"),
)
print(RESULT_MARKER_PREFIX + base64.b64encode(payload.encode("utf-8")).decode("ascii"))
if __name__ == "__main__":
args = json.loads(sys.argv[1])
result = main(**args)
if result is not None:
print(result)
emit_result(result)
""")
elif language == SupportLanguage.NODEJS:
code_name = "main.js"
code_path = os.path.join(workdir, "main.js")
code_path = os.path.join(workdir, code_name)
with open(code_path, "wb") as f:
f.write(base64.b64decode(req.code_b64))
runner_name = "runner.js"
runner_path = os.path.join(workdir, "runner.js")
with open(runner_path, "w") as f:
f.write("""
runner_code = """
const fs = require('fs');
const path = require('path');
const args = JSON.parse(process.argv[2]);
const mainPath = path.join(__dirname, 'main.js');
const RESULT_MARKER_PREFIX = '__RESULT_MARKER_PREFIX__';
function isPromise(value) {
return Boolean(value && typeof value.then === 'function');
}
function emitResult(value) {
if (typeof value === 'undefined') {
console.error('Error: main() must return a value. Use null for an empty result.');
process.exit(1);
}
const payload = JSON.stringify({ present: true, value, type: 'json' });
if (typeof payload === 'undefined') {
console.error('Error: main() returned a non-JSON-serializable value.');
process.exit(1);
}
console.log(RESULT_MARKER_PREFIX + Buffer.from(payload, 'utf8').toString('base64'));
}
if (fs.existsSync(mainPath)) {
const mod = require(mainPath);
const main = typeof mod === 'function' ? mod : mod.main;
@@ -103,40 +161,38 @@ if (fs.existsSync(mainPath)) {
if (typeof args === 'object' && args !== null) {
try {
const result = main(args);
const result = Promise.resolve(main(args));
if (isPromise(result)) {
result.then(output => {
if (output !== null) {
console.log(output);
}
emitResult(output);
}).catch(err => {
console.error('Error in async main function:', err);
process.exit(1);
});
} else {
if (result !== null) {
console.log(result);
}
emitResult(result);
}
} catch (err) {
console.error('Error when executing main:', err);
process.exit(1);
}
} else {
console.error('Error: args is not a valid object:', args);
process.exit(1);
}
} else {
console.error('main.js not found in the current directory');
process.exit(1);
}
""")
# dirs
"""
f.write(runner_code.replace("__RESULT_MARKER_PREFIX__", RESULT_MARKER_PREFIX))
returncode, _, stderr = await async_run_command("docker", "exec", container, "mkdir", "-p", f"/workspace/{task_id}", timeout=5)
if returncode != 0:
raise RuntimeError(f"Directory creation failed: {stderr}")
# archive
tar_proc = await asyncio.create_subprocess_exec("tar", "czf", "-", "-C", workdir, code_name, runner_name, stdout=asyncio.subprocess.PIPE)
tar_stdout, _ = await tar_proc.communicate()
# unarchive
docker_proc = await asyncio.create_subprocess_exec(
"docker", "exec", "-i", container, "tar", "xzf", "-", "-C", f"/workspace/{task_id}", stdin=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
)
@@ -145,7 +201,6 @@ if (fs.existsSync(mainPath)) {
if docker_proc.returncode != 0:
raise RuntimeError(stderr.decode())
# exec
start_time = time.time()
try:
logger.info(f"Passed in args: {req.arguments}")
@@ -160,11 +215,10 @@ if (fs.existsSync(mainPath)) {
str(TIMEOUT),
language,
]
# flags
if language == SupportLanguage.PYTHON:
run_args.extend(["-I", "-B"])
elif language == SupportLanguage.NODEJS:
run_args.extend([])
pass # no additional flags
else:
assert False, "Will never reach here"
run_args.extend([runner_name, args_json])
@@ -184,14 +238,16 @@ if (fs.existsSync(mainPath)) {
logger.info(f"{args_json=}")
if returncode == 0:
clean_stdout, structured_result = _extract_result_envelope(stdout)
artifacts = await _collect_artifacts(container, task_id, workdir)
return CodeExecutionResult(
status=ResultStatus.SUCCESS,
stdout=str(stdout),
stdout=clean_stdout,
stderr=stderr,
exit_code=0,
time_used_ms=time_used_ms,
artifacts=artifacts,
result=structured_result,
)
elif returncode == 124:
return CodeExecutionResult(
@@ -229,7 +285,6 @@ if (fs.existsSync(mainPath)) {
return CodeExecutionResult(status=ResultStatus.PROGRAM_RUNNER_ERROR, stdout="", stderr=str(e), exit_code=-3, detail="internal_error")
finally:
# cleanup
cleanup_tasks = [async_run_command("docker", "exec", container, "rm", "-rf", f"/workspace/{task_id}"), async_run_command("rm", "-rf", workdir)]
await asyncio.gather(*cleanup_tasks, return_exceptions=True)
await release_container(container, language)

30
agent/sandbox/executor_manager/services/security.py
View File

@@ -14,6 +14,7 @@
# limitations under the License.
#
import ast
import re
from typing import List, Tuple
from core.logger import logger
@@ -151,6 +152,26 @@ class SecurePythonAnalyzer(ast.NodeVisitor):
self.generic_visit(node)
class SecureJavaScriptAnalyzer:
DANGEROUS_PATTERNS = [
(re.compile(r"""require\s*\(\s*['"]child_process['"]\s*\)"""), "Require: child_process"),
(re.compile(r"""require\s*\(\s*['"]fs['"]\s*\)"""), "Require: fs"),
(re.compile(r"""require\s*\(\s*['"]worker_threads['"]\s*\)"""), "Require: worker_threads"),
(re.compile(r"""\beval\s*\("""), "Call: eval"),
(re.compile(r"""\bFunction\s*\("""), "Call: Function"),
(re.compile(r"""\bprocess\s*\.\s*binding\s*\("""), "Call: process.binding"),
]
@classmethod
def analyze(cls, code: str) -> List[Tuple[str, int]]:
issues: List[Tuple[str, int]] = []
for pattern, description in cls.DANGEROUS_PATTERNS:
for match in pattern.finditer(code):
lineno = code.count("\n", 0, match.start()) + 1
issues.append((description, lineno))
return issues
def analyze_code_security(code: str, language: SupportLanguage) -> Tuple[bool, List[Tuple[str, int]]]:
"""
Analyze the provided code string and return whether it's safe and why.
@@ -168,6 +189,9 @@ def analyze_code_security(code: str, language: SupportLanguage) -> Tuple[bool, L
except Exception as e:
logger.error(f"[SafeCheck] Python parsing failed: {str(e)}")
return False, [(f"Parsing Error: {str(e)}", -1)]
else:
logger.warning(f"[SafeCheck] Unsupported language for security analysis: {language} — defaulting to SAFE (manual review recommended)")
return True, [(f"Unsupported language for security analysis: {language} — defaulted to SAFE, manual review recommended", -1)]
if language == SupportLanguage.NODEJS:
issues = SecureJavaScriptAnalyzer.analyze(code)
return len(issues) == 0, issues
logger.warning(f"[SafeCheck] Unsupported language for security analysis: {language}")
return False, [(f"Unsupported language for security analysis: {language}", -1)]

127
agent/sandbox/providers/aliyun_codeinterpreter.py
View File

@@ -30,6 +30,8 @@ https://api.aliyun.com/api/AgentRun/2025-09-10/CreateSandbox?lang=PYTHON
import logging
import os
import time
import base64
import json
from typing import Dict, Any, List, Optional
from datetime import datetime, timezone
@@ -40,6 +42,7 @@ from agentrun.utils.exception import ServerError
from .base import SandboxProvider, SandboxInstance, ExecutionResult
logger = logging.getLogger(__name__)
RESULT_MARKER_PREFIX = "__RAGFLOW_RESULT__:"
class AliyunCodeInterpreterProvider(SandboxProvider):
@@ -51,9 +54,9 @@ class AliyunCodeInterpreterProvider(SandboxProvider):
"""
def __init__(self):
self.access_key_id: Optional[str] = None
self.access_key_secret: Optional[str] = None
self.account_id: Optional[str] = None
self.access_key_id: Optional[str] = ""
self.access_key_secret: Optional[str] = ""
self.account_id: Optional[str] = ""
self.region: str = "cn-hangzhou"
self.template_name: str = ""
self.timeout: int = 30
@@ -146,8 +149,6 @@ class AliyunCodeInterpreterProvider(SandboxProvider):
try:
# Get or create template
from agentrun.sandbox import Sandbox
if self.template_name:
# Use existing template
template_name = self.template_name
@@ -226,48 +227,17 @@ class AliyunCodeInterpreterProvider(SandboxProvider):
# Connect to existing sandbox instance
sandbox = Sandbox.connect(sandbox_id=instance_id, config=self._config)
# Convert language string to CodeLanguage enum
code_language = CodeLanguage.PYTHON if normalized_lang == "python" else CodeLanguage.JAVASCRIPT
# agentrun-sdk 0.0.26 only exposes CodeLanguage.PYTHON; keep JS as string fallback.
code_language = CodeLanguage.PYTHON if normalized_lang == "python" else "javascript"
# Wrap code to call main() function
# Matches self_managed provider behavior: call main(**arguments)
if normalized_lang == "python":
# Build arguments string for main() call
if arguments:
import json as json_module
args_json = json_module.dumps(arguments)
wrapped_code = f'''{code}
if __name__ == "__main__":
import json
result = main(**{args_json})
print(json.dumps(result) if isinstance(result, dict) else result)
'''
else:
wrapped_code = f'''{code}
if __name__ == "__main__":
import json
result = main()
print(json.dumps(result) if isinstance(result, dict) else result)
'''
else: # javascript
if arguments:
import json as json_module
args_json = json_module.dumps(arguments)
wrapped_code = f'''{code}
// Call main and output result
const result = main({args_json});
console.log(typeof result === 'object' ? JSON.stringify(result) : String(result));
'''
else:
wrapped_code = f'''{code}
// Call main and output result
const result = main();
console.log(typeof result === 'object' ? JSON.stringify(result) : String(result));
'''
args_json = json.dumps(arguments or {})
wrapped_code = (
self._build_python_wrapper(code, args_json)
if normalized_lang == "python"
else self._build_javascript_wrapper(code, args_json)
)
logger.debug(f"Aliyun Code Interpreter: Wrapped code (first 200 chars): {wrapped_code[:200]}")
start_time = time.time()
@@ -314,6 +284,7 @@ console.log(typeof result === 'object' ? JSON.stringify(result) : String(result)
stdout = "\n".join(stdout_parts)
stderr = "\n".join(stderr_parts)
stdout, structured_result = self._extract_structured_result(stdout)
logger.info(f"Aliyun Code Interpreter: stdout length={len(stdout)}, stderr length={len(stderr)}, exit_code={exit_code}")
if stdout:
@@ -331,6 +302,9 @@ console.log(typeof result === 'object' ? JSON.stringify(result) : String(result)
"language": normalized_lang,
"context_id": result.get("contextId") if isinstance(result, dict) else None,
"timeout": timeout,
"result_present": structured_result.get("present", False),
"result_value": structured_result.get("value"),
"result_type": structured_result.get("type"),
},
)
@@ -390,6 +364,71 @@ console.log(typeof result === 'object' ? JSON.stringify(result) : String(result)
# If we get any response (even an error), the service is reachable
return "connection" not in str(e).lower()
@staticmethod
def _build_python_wrapper(code: str, args_json: str) -> str:
marker = RESULT_MARKER_PREFIX
return f'''{code}
if __name__ == "__main__":
import base64
import json
result = main(**{args_json})
payload = json.dumps({{"present": True, "value": result, "type": "json"}}, ensure_ascii=False, separators=(",", ":"))
print("{marker}" + base64.b64encode(payload.encode("utf-8")).decode("ascii"))
'''
@staticmethod
def _build_javascript_wrapper(code: str, args_json: str) -> str:
marker = RESULT_MARKER_PREFIX
return f'''{code}
const __ragflowArgs = {args_json};
(async () => {{
try {{
const output = await Promise.resolve(main(__ragflowArgs));
if (typeof output === 'undefined') {{
throw new Error('main() must return a value. Use null for an empty result.');
}}
const payload = JSON.stringify({{ present: true, value: output, type: 'json' }});
if (typeof payload === 'undefined') {{
throw new Error('main() returned a non-JSON-serializable value.');
}}
console.log('{marker}' + Buffer.from(payload, 'utf8').toString('base64'));
}} catch (err) {{
console.error(err instanceof Error ? err.stack || err.message : String(err));
}}
}})();
'''
@staticmethod
def _extract_structured_result(stdout: str) -> tuple[str, Dict[str, Any]]:
if not stdout:
return "", {}
cleaned_lines: list[str] = []
structured_result: Dict[str, Any] = {}
for line in str(stdout).splitlines():
if line.startswith(RESULT_MARKER_PREFIX):
payload_b64 = line[len(RESULT_MARKER_PREFIX) :].strip()
if not payload_b64:
continue
try:
payload = base64.b64decode(payload_b64).decode("utf-8")
structured_result = json.loads(payload)
except Exception as exc:
logger.warning(f"Aliyun Code Interpreter: failed to decode structured result marker: {exc}")
cleaned_lines.append(line)
continue
cleaned_lines.append(line)
cleaned_stdout = "\n".join(cleaned_lines)
if stdout.endswith("\n") and cleaned_stdout and not cleaned_stdout.endswith("\n"):
cleaned_stdout += "\n"
return cleaned_stdout, structured_result
def get_supported_languages(self) -> List[str]:
"""
Get list of supported programming languages.

4
agent/sandbox/providers/self_managed.py
View File

@@ -187,6 +187,7 @@ class SelfManagedProvider(SandboxProvider):
)
result = response.json()
structured_result = result.get("result") or {}
return ExecutionResult(
stdout=result.get("stdout", ""),
@@ -200,6 +201,9 @@ class SelfManagedProvider(SandboxProvider):
"detail": result.get("detail"),
"instance_id": instance_id,
"artifacts": result.get("artifacts", []),
"result_present": structured_result.get("present", False),
"result_value": structured_result.get("value"),
"result_type": structured_result.get("type"),
}
)

56
agent/sandbox/tests/test_aliyun_codeinterpreter.py
View File

@@ -101,13 +101,15 @@ class TestAliyunCodeInterpreterProvider:
assert provider.region == "cn-hangzhou"
assert provider.template_name == ""
@patch("agent.sandbox.providers.aliyun_codeinterpreter.CodeInterpreterSandbox")
def test_create_instance_python(self, mock_sandbox_class):
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Template")
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Sandbox")
def test_create_instance_python(self, mock_sandbox_class, mock_template):
"""Test creating a Python instance."""
# Mock successful instance creation
mock_sandbox = MagicMock()
mock_sandbox.sandbox_id = "01JCED8Z9Y6XQVK8M2NRST5WXY"
mock_sandbox_class.return_value = mock_sandbox
mock_sandbox_class.create.return_value = mock_sandbox
mock_template.get_by_name.return_value = MagicMock()
provider = AliyunCodeInterpreterProvider()
provider._initialized = True
@@ -119,12 +121,14 @@ class TestAliyunCodeInterpreterProvider:
assert instance.status == "READY"
assert instance.metadata["language"] == "python"
@patch("agent.sandbox.providers.aliyun_codeinterpreter.CodeInterpreterSandbox")
def test_create_instance_javascript(self, mock_sandbox_class):
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Template")
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Sandbox")
def test_create_instance_javascript(self, mock_sandbox_class, mock_template):
"""Test creating a JavaScript instance."""
mock_sandbox = MagicMock()
mock_sandbox.sandbox_id = "01JCED8Z9Y6XQVK8M2NRST5WXY"
mock_sandbox_class.return_value = mock_sandbox
mock_sandbox_class.create.return_value = mock_sandbox
mock_template.get_by_name.return_value = MagicMock()
provider = AliyunCodeInterpreterProvider()
provider._initialized = True
@@ -141,7 +145,7 @@ class TestAliyunCodeInterpreterProvider:
with pytest.raises(RuntimeError, match="Provider not initialized"):
provider.create_instance("python")
@patch("agent.sandbox.providers.aliyun_codeinterpreter.CodeInterpreterSandbox")
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Sandbox")
def test_execute_code_success(self, mock_sandbox_class):
"""Test successful code execution."""
# Mock sandbox instance
@@ -150,7 +154,7 @@ class TestAliyunCodeInterpreterProvider:
"results": [{"type": "stdout", "text": "Hello, World!"}, {"type": "result", "text": "None"}, {"type": "endOfExecution", "status": "ok"}],
"contextId": "kernel-12345-67890",
}
mock_sandbox_class.return_value = mock_sandbox
mock_sandbox_class.connect.return_value = mock_sandbox
provider = AliyunCodeInterpreterProvider()
provider._initialized = True
@@ -163,14 +167,14 @@ class TestAliyunCodeInterpreterProvider:
assert result.exit_code == 0
assert result.execution_time > 0
@patch("agent.sandbox.providers.aliyun_codeinterpreter.CodeInterpreterSandbox")
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Sandbox")
def test_execute_code_timeout(self, mock_sandbox_class):
"""Test code execution timeout."""
from agentrun.utils.exception import ServerError
mock_sandbox = MagicMock()
mock_sandbox.context.execute.side_effect = ServerError(408, "Request timeout")
mock_sandbox_class.return_value = mock_sandbox
mock_sandbox_class.connect.return_value = mock_sandbox
provider = AliyunCodeInterpreterProvider()
provider._initialized = True
@@ -179,14 +183,14 @@ class TestAliyunCodeInterpreterProvider:
with pytest.raises(TimeoutError, match="Execution timed out"):
provider.execute_code(instance_id="01JCED8Z9Y6XQVK8M2NRST5WXY", code="while True: pass", language="python", timeout=5)
@patch("agent.sandbox.providers.aliyun_codeinterpreter.CodeInterpreterSandbox")
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Sandbox")
def test_execute_code_with_error(self, mock_sandbox_class):
"""Test code execution with error."""
mock_sandbox = MagicMock()
mock_sandbox.context.execute.return_value = {
"results": [{"type": "stderr", "text": "Traceback..."}, {"type": "error", "text": "NameError: name 'x' is not defined"}, {"type": "endOfExecution", "status": "error"}]
}
mock_sandbox_class.return_value = mock_sandbox
mock_sandbox_class.connect.return_value = mock_sandbox
provider = AliyunCodeInterpreterProvider()
provider._initialized = True
@@ -197,6 +201,34 @@ class TestAliyunCodeInterpreterProvider:
assert result.exit_code != 0
assert len(result.stderr) > 0
@patch("agent.sandbox.providers.aliyun_codeinterpreter.Sandbox")
def test_execute_code_uses_structured_result_marker_for_async_javascript(self, mock_sandbox_class):
"""Test JavaScript wrapper uses the structured result marker and awaits async main."""
mock_sandbox = MagicMock()
mock_sandbox.context.execute.return_value = {
"results": [{"type": "stdout", "text": "__RAGFLOW_RESULT__:eyJwcmVzZW50Ijp0cnVlLCJ2YWx1ZSI6eyJhIjoiYiJ9LCJ0eXBlIjoianNvbiJ9"}],
"contextId": "kernel-12345-67890",
}
mock_sandbox_class.connect.return_value = mock_sandbox
provider = AliyunCodeInterpreterProvider()
provider._initialized = True
provider._config = MagicMock()
result = provider.execute_code(
instance_id="01JCED8Z9Y6XQVK8M2NRST5WXY",
code="async function main(args) { return { a: 'b' }; }",
language="javascript",
timeout=10,
)
wrapped_code = mock_sandbox.context.execute.call_args.kwargs["code"]
assert "__RAGFLOW_RESULT__:" in wrapped_code
assert "await Promise.resolve(main(" in wrapped_code
assert result.metadata["result_present"] is True
assert result.metadata["result_value"] == {"a": "b"}
assert result.metadata["result_type"] == "json"
def test_get_supported_languages(self):
"""Test getting supported languages."""
provider = AliyunCodeInterpreterProvider()

35
agent/sandbox/tests/test_providers.py
View File

@@ -254,6 +254,41 @@ class TestSelfManagedProvider:
assert result.metadata["status"] == "success"
assert result.metadata["instance_id"] == "test-123"
@patch('requests.post')
def test_execute_code_maps_structured_result_into_metadata(self, mock_post):
"""Test successful code execution with structured result envelope."""
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {
"status": "success",
"stdout": "debug line\n",
"stderr": "",
"exit_code": 0,
"time_used_ms": 100.0,
"memory_used_kb": 1024.0,
"result": {
"present": True,
"value": {"items": ["a", "b"]},
"type": "json",
},
}
mock_post.return_value = mock_response
provider = SelfManagedProvider()
provider._initialized = True
result = provider.execute_code(
instance_id="test-123",
code="def main(): return {'items': ['a', 'b']}",
language="python",
timeout=10
)
assert result.stdout == "debug line\n"
assert result.metadata["result_present"] is True
assert result.metadata["result_value"] == {"items": ["a", "b"]}
assert result.metadata["result_type"] == "json"
@patch('requests.post')
def test_execute_code_timeout(self, mock_post):
"""Test code execution timeout."""

55
agent/sandbox/tests/test_security.py Normal file
View File

@@ -0,0 +1,55 @@
#
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import sys
from pathlib import Path
EXECUTOR_MANAGER_ROOT = Path(__file__).resolve().parents[1] / "executor_manager"
if str(EXECUTOR_MANAGER_ROOT) not in sys.path:
sys.path.insert(0, str(EXECUTOR_MANAGER_ROOT))
from models.enums import SupportLanguage # noqa: E402
from services.security import analyze_code_security # noqa: E402
def test_javascript_child_process_is_rejected():
is_safe, issues = analyze_code_security(
"const cp = require('child_process'); async function main() { return 'ok'; }",
SupportLanguage.NODEJS,
)
assert is_safe is False
assert any("child_process" in issue for issue, _ in issues)
def test_javascript_eval_is_rejected():
is_safe, issues = analyze_code_security(
"async function main() { return eval('1+1'); }",
SupportLanguage.NODEJS,
)
assert is_safe is False
assert any("eval" in issue.lower() for issue, _ in issues)
def test_javascript_safe_code_still_passes():
is_safe, issues = analyze_code_security(
"async function main(args) { return { answer: args.value ?? null }; }",
SupportLanguage.NODEJS,
)
assert is_safe is True
assert issues == []

413
agent/tools/code_exec.py
View File

@@ -20,6 +20,7 @@ import logging
import os
import uuid
from abc import ABC
from collections.abc import Mapping
from typing import Optional
from pydantic import BaseModel, Field, field_validator
@@ -32,6 +33,183 @@ from common.connection_utils import timeout
from common.constants import SANDBOX_ARTIFACT_BUCKET, SANDBOX_ARTIFACT_EXPIRE_DAYS
SYSTEM_OUTPUT_KEYS = frozenset(
{
"content",
"actual_type",
"_ERROR",
"_ARTIFACTS",
"_ATTACHMENT_CONTENT",
"raw_result",
"_created_time",
"_elapsed_time",
}
)
class ContractError(ValueError):
pass
def _validate_business_output_name(name: str) -> None:
if not name or not name.strip():
raise ContractError("CodeExec business output name must not be empty")
if name in SYSTEM_OUTPUT_KEYS:
raise ContractError(f"CodeExec reserved output name is not allowed: {name}")
if "." in name:
raise ContractError(f"CodeExec business output name must not contain '.': {name}")
def select_business_output(outputs: Mapping[str, object]) -> tuple[str, object]:
if len(outputs) == 1:
only_name, only_meta = next(iter(outputs.items()))
_validate_business_output_name(only_name)
return only_name, only_meta
business_outputs = [(name, meta) for name, meta in outputs.items() if name not in SYSTEM_OUTPUT_KEYS]
if len(business_outputs) != 1:
raise ContractError(
f"CodeExec contract must contain exactly one business output, got {len(business_outputs)}"
)
_validate_business_output_name(business_outputs[0][0])
return business_outputs[0]
def normalize_output_value(value):
if isinstance(value, (tuple, list)):
return [normalize_output_value(item) for item in value]
if isinstance(value, dict):
return {key: normalize_output_value(item) for key, item in value.items()}
return value
def infer_actual_type(value) -> str:
value = normalize_output_value(value)
if value is None:
return "Null"
if isinstance(value, bool):
return "Boolean"
if _is_number(value):
return "Number"
if isinstance(value, str):
return "String"
if isinstance(value, dict):
return "Object"
if isinstance(value, list):
if not value:
return "Array<Any>"
inferred = {infer_actual_type(item) for item in value}
if len(inferred) == 1:
return f"Array<{inferred.pop()}>"
return "Array<Any>"
return "Any"
def render_canonical_content(value) -> str:
value = normalize_output_value(value)
if value is None:
return ""
if isinstance(value, str):
return value
if isinstance(value, (dict, list)):
return json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True)
return str(value)
def _is_number(value) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool)
def _validate_top_level_value_domain(value) -> None:
allowed = value is None or isinstance(value, (bool, str, dict, list)) or _is_number(value)
if not allowed:
raise ContractError(
f"CodeExec unsupported top-level result type: {type(value).__name__}. "
"Allowed top-level values are String, Number, Boolean, Object, Array, or Null."
)
def _normalize_expected_type(expected_type: str) -> str:
etype = expected_type.strip()
low = etype.lower()
simple_types = {
"string": "String",
"number": "Number",
"boolean": "Boolean",
"object": "Object",
"null": "Null",
"any": "Any",
}
if low in simple_types:
return simple_types[low]
if low.startswith("array<") and low.endswith(">"):
inner = etype[etype.find("<") + 1 : -1].strip()
if not inner:
raise ContractError(f"Unsupported expected type: {expected_type}")
return f"Array<{_normalize_expected_type(inner)}>"
return etype
def _validate_expected_type(expected_type: str, value, path: str = "") -> None:
etype = _normalize_expected_type(expected_type)
if not etype or etype.lower() == "any":
return
value = normalize_output_value(value)
if etype.startswith("Array<") and etype.endswith(">"):
inner_type = etype[6:-1].strip()
if not isinstance(value, list):
raise ContractError(
f"CodeExec contract mismatch at {path or 'value'}: expected type {etype}, got {infer_actual_type(value)}"
)
for index, item in enumerate(value):
child_path = f"{path}[{index}]" if path else f"[{index}]"
_validate_expected_type(inner_type, item, child_path)
return
actual_type = infer_actual_type(value)
if etype == "String":
valid = isinstance(value, str)
elif etype == "Number":
valid = _is_number(value)
elif etype == "Boolean":
valid = isinstance(value, bool)
elif etype == "Object":
valid = isinstance(value, dict)
elif etype == "Null":
valid = value is None
else:
raise ContractError(f"Unsupported expected type: {expected_type}")
if not valid:
raise ContractError(
f"CodeExec contract mismatch at {path or 'value'}: expected type {etype}, got {actual_type}"
)
def build_code_exec_contract(outputs: Mapping[str, object], raw_result) -> dict[str, object]:
business_name, business_meta = select_business_output(outputs)
expected_type = ""
if isinstance(business_meta, Mapping):
expected_type = str(business_meta.get("type") or "")
normalized_value = normalize_output_value(raw_result)
_validate_top_level_value_domain(normalized_value)
_validate_expected_type(expected_type, normalized_value)
return {
"business_output": business_name,
"value": normalized_value,
"actual_type": infer_actual_type(normalized_value),
"content": render_canonical_content(normalized_value),
}
def _art_field(art, field: str, default=""):
return art.get(field, default) if isinstance(art, dict) else getattr(art, field, default)
class Language(StrEnum):
PYTHON = "python"
NODEJS = "nodejs"
@@ -190,7 +368,13 @@ class CodeExec(ToolBase, ABC):
return
artifacts = result.metadata.get("artifacts", []) if result.metadata else []
return self._process_execution_result(result.stdout, result.stderr, "Provider system", artifacts)
return self._process_execution_result(
result.stdout,
result.stderr,
"Provider system",
artifacts,
execution_metadata=result.metadata,
)
except (ImportError, RuntimeError) as provider_error:
# Provider system not available or not configured, fall back to HTTP
@@ -226,6 +410,7 @@ class CodeExec(ToolBase, ABC):
body.get("stderr"),
f"http://{settings.SANDBOX_HOST}:9385/run",
body.get("artifacts", []),
execution_metadata=self._build_http_execution_metadata(body),
)
else:
self.set_output("_ERROR", "There is no response from sandbox")
@@ -239,8 +424,18 @@ class CodeExec(ToolBase, ABC):
return self.output()
def _process_execution_result(self, stdout: str, stderr: str | None, source: str, artifacts: list | None = None):
if stderr and not stdout and not artifacts:
def _process_execution_result(
self,
stdout: str,
stderr: str | None,
source: str,
artifacts: list | None = None,
execution_metadata: dict | None = None,
):
has_structured_result = bool((execution_metadata or {}).get("result_present") is True)
resolved_value, used_stdout_fallback = self._resolve_execution_result_value(stdout, execution_metadata)
if stderr and not has_structured_result and not artifacts and not str(stdout or "").strip():
self.set_output("_ERROR", stderr)
return self.output()
@@ -250,29 +445,48 @@ class CodeExec(ToolBase, ABC):
if stderr:
logging.warning(f"[CodeExec]: stderr (non-fatal): {stderr[:500]}")
parsed_stdout = self._deserialize_stdout(stdout)
logging.info(f"[CodeExec]: {source} -> {parsed_stdout}")
self._populate_outputs(parsed_stdout, stdout)
if used_stdout_fallback and str(stdout or "").strip():
logging.warning("[CodeExec]: Falling back to stdout deserialization because no structured result metadata was provided")
logging.info(f"[CodeExec]: {source} -> {resolved_value}")
content_parts = []
base_content = self._build_content_text(parsed_stdout, raw_stdout=stdout)
base_content = self._apply_business_output(resolved_value)
if base_content:
content_parts.append(base_content)
if artifacts:
artifact_urls = self._upload_artifacts(artifacts)
if artifact_urls:
self.set_output("_ARTIFACTS", artifact_urls)
self.set_output("_ARTIFACTS", artifact_urls or None)
attachment_text = self._build_attachment_content(artifacts, artifact_urls)
self.set_output("_ATTACHMENT_CONTENT", attachment_text)
if attachment_text:
content_parts.append(attachment_text)
else:
self.set_output("_ARTIFACTS", None)
self.set_output("_ATTACHMENT_CONTENT", "")
self.set_output("content", "\n\n".join([part for part in content_parts if part]).strip())
return self.output()
def _build_http_execution_metadata(self, body: Mapping | None) -> dict:
if not isinstance(body, Mapping):
return {}
structured_result = body.get("result")
if not isinstance(structured_result, Mapping):
return {}
return {
"result_present": structured_result.get("present", False),
"result_value": structured_result.get("value"),
"result_type": structured_result.get("type"),
}
def _resolve_execution_result_value(self, stdout: str, execution_metadata: Mapping | None = None):
metadata = execution_metadata or {}
if metadata.get("result_present") is True:
return metadata.get("result_value"), False
return self._deserialize_stdout(stdout), True
@classmethod
def _ensure_bucket_lifecycle(cls):
if cls._lifecycle_configured:
@@ -306,10 +520,10 @@ class CodeExec(ToolBase, ABC):
uploaded = []
for art in artifacts:
try:
name = art.get("name", "") if isinstance(art, dict) else getattr(art, "name", "")
content_b64 = art.get("content_b64", "") if isinstance(art, dict) else getattr(art, "content_b64", "")
mime_type = art.get("mime_type", "") if isinstance(art, dict) else getattr(art, "mime_type", "")
size = art.get("size", 0) if isinstance(art, dict) else getattr(art, "size", 0)
name = _art_field(art, "name")
content_b64 = _art_field(art, "content_b64")
mime_type = _art_field(art, "mime_type")
size = _art_field(art, "size", 0)
if not content_b64 or not name:
continue
@@ -350,119 +564,24 @@ class CodeExec(ToolBase, ABC):
continue
return text
def _coerce_output_value(self, value, expected_type: Optional[str]):
if expected_type is None:
return value
etype = expected_type.strip().lower()
inner_type = None
if etype.startswith("array<") and etype.endswith(">"):
inner_type = etype[6:-1].strip()
etype = "array"
def _apply_business_output(self, parsed_stdout) -> str:
normalized_result = normalize_output_value(parsed_stdout)
self.set_output("raw_result", normalized_result)
business_output_names = [name for name in self._param.outputs if name not in SYSTEM_OUTPUT_KEYS]
try:
if etype == "string":
return "" if value is None else str(value)
contract = build_code_exec_contract(self._param.outputs, normalized_result)
except ContractError as e:
for output_name in business_output_names:
self.set_output(output_name, None)
self.set_output("actual_type", infer_actual_type(normalized_result))
self.set_output("_ERROR", str(e))
logging.warning(f"[CodeExec]: contract validation failed: {e}")
return render_canonical_content(normalized_result)
if etype == "number":
if value is None or value == "":
return None
if isinstance(value, (int, float)):
return value
if isinstance(value, str):
try:
return float(value)
except Exception:
return value
return float(value)
if etype == "boolean":
if isinstance(value, bool):
return value
if isinstance(value, str):
lv = value.lower()
if lv in ("true", "1", "yes", "y", "on"):
return True
if lv in ("false", "0", "no", "n", "off"):
return False
return bool(value)
if etype == "array":
candidate = value
if isinstance(candidate, str):
parsed = self._deserialize_stdout(candidate)
candidate = parsed
if isinstance(candidate, tuple):
candidate = list(candidate)
if not isinstance(candidate, list):
candidate = [] if candidate is None else [candidate]
if inner_type == "string":
return ["" if v is None else str(v) for v in candidate]
if inner_type == "number":
coerced = []
for v in candidate:
try:
if v is None or v == "":
coerced.append(None)
elif isinstance(v, (int, float)):
coerced.append(v)
else:
coerced.append(float(v))
except Exception:
coerced.append(v)
return coerced
return candidate
if etype == "object":
if isinstance(value, dict):
return value
if isinstance(value, str):
parsed = self._deserialize_stdout(value)
if isinstance(parsed, dict):
return parsed
return value
except Exception:
return value
return value
def _populate_outputs(self, parsed_stdout, raw_stdout: str):
outputs_items = list(self._param.outputs.items())
logging.info(f"[CodeExec]: outputs schema keys: {[k for k, _ in outputs_items]}")
if not outputs_items:
return
if isinstance(parsed_stdout, dict):
for key, meta in outputs_items:
if key.startswith("_"):
continue
val = self._get_by_path(parsed_stdout, key)
if val is None and len(outputs_items) == 1:
val = parsed_stdout
coerced = self._coerce_output_value(val, meta.get("type"))
logging.info(f"[CodeExec]: populate dict key='{key}' raw='{val}' coerced='{coerced}'")
self.set_output(key, coerced)
return
if isinstance(parsed_stdout, (list, tuple)):
for idx, (key, meta) in enumerate(outputs_items):
if key.startswith("_"):
continue
val = parsed_stdout[idx] if idx < len(parsed_stdout) else None
coerced = self._coerce_output_value(val, meta.get("type"))
logging.info(f"[CodeExec]: populate list key='{key}' raw='{val}' coerced='{coerced}'")
self.set_output(key, coerced)
return
default_val = parsed_stdout if parsed_stdout is not None else raw_stdout
for idx, (key, meta) in enumerate(outputs_items):
if key.startswith("_"):
continue
val = default_val if idx == 0 else None
coerced = self._coerce_output_value(val, meta.get("type"))
logging.info(f"[CodeExec]: populate scalar key='{key}' raw='{val}' coerced='{coerced}'")
self.set_output(key, coerced)
self.set_output("actual_type", contract["actual_type"])
self.set_output(contract["business_output"], contract["value"])
return contract["content"]
def _build_attachment_content(self, artifacts: list, artifact_urls: list[dict] | None = None) -> str:
sections = []
@@ -471,9 +590,9 @@ class CodeExec(ToolBase, ABC):
for idx, art in enumerate(artifacts, start=1):
key = f"attachment{idx}"
try:
name = art.get("name", "") if isinstance(art, dict) else getattr(art, "name", "")
content_b64 = art.get("content_b64", "") if isinstance(art, dict) else getattr(art, "content_b64", "")
mime_type = art.get("mime_type", "") if isinstance(art, dict) else getattr(art, "mime_type", "")
name = _art_field(art, "name")
content_b64 = _art_field(art, "content_b64")
mime_type = _art_field(art, "mime_type")
if not name or not content_b64:
continue
@@ -490,11 +609,8 @@ class CodeExec(ToolBase, ABC):
logging.info(f"[CodeExec]: parse attachment section key='{key}' from artifact='{name}'")
except Exception as e:
logging.warning(f"[CodeExec]: Failed to parse artifact for content section '{key}': {e}")
fallback_type = self._normalize_attachment_type(
art.get("name", "") if isinstance(art, dict) else getattr(art, "name", ""),
art.get("mime_type", "") if isinstance(art, dict) else getattr(art, "mime_type", ""),
)
fallback_name = art.get("name", "") if isinstance(art, dict) else getattr(art, "name", "")
fallback_type = self._normalize_attachment_type(name, mime_type)
fallback_name = name
fallback_url = ""
if idx - 1 < len(artifact_urls):
fallback_url = artifact_urls[idx - 1].get("url", "")
@@ -529,38 +645,3 @@ class CodeExec(ToolBase, ABC):
title += f": {name}"
body = parsed if isinstance(parsed, str) else json.dumps(parsed, ensure_ascii=False)
return f"{title}\n{body}".strip()
def _build_content_text(self, parsed_stdout, raw_stdout: str = "") -> str:
if isinstance(parsed_stdout, str):
return parsed_stdout.strip()
if isinstance(parsed_stdout, (dict, list, tuple)):
try:
return json.dumps(parsed_stdout, ensure_ascii=False, indent=2).strip()
except Exception:
return str(parsed_stdout).strip()
if parsed_stdout is None:
return str(raw_stdout or "").strip()
return str(parsed_stdout).strip()
def _get_by_path(self, data, path: str):
if not path:
return None
cur = data
for part in path.split("."):
part = part.strip()
if not part:
return None
if isinstance(cur, dict):
cur = cur.get(part)
elif isinstance(cur, list):
try:
idx = int(part)
cur = cur[idx]
except Exception:
return None
else:
return None
if cur is None:
return None
logging.info(f"[CodeExec]: resolve path '{path}' -> {cur}")
return cur