refactor(sandbox): extract connection helpers and move run_command to helper module

- Add helpers.py with connection management utilities:
    - with_connection: context manager for connection lifecycle
    - submit_command: execute command and return CommandFuture
    - execute: run command with auto connection, raise on failure
    - try_execute: run command with auto connection, return result

  - Add CommandExecutionError to exec.py for typed error handling
    with access to exit_code, stderr, and full result

  - Remove run_command method from VirtualEnvironment base class
    (now available as submit_command helper)

  - Update all call sites to use new helper functions:
    - sandbox/session.py
    - sandbox/storage/archive_storage.py
    - sandbox/bash/bash_tool.py
    - workflow/nodes/command/node.py

  - Add comprehensive unit tests for helpers with connection reuse
This commit is contained in:
Harry
2026-01-14 23:23:00 +08:00
parent 31427e9c42
commit a0c388f283
19 changed files with 553 additions and 149 deletions

View File

@ -1,12 +1,12 @@
import contextlib
import logging
import shlex
from collections.abc import Mapping, Sequence
from typing import Any
from core.sandbox.debug import sandbox_debug
from core.sandbox.manager import SandboxManager
from core.sandbox.utils.debug import sandbox_debug
from core.virtual_environment.__base.command_future import CommandCancelledError, CommandTimeoutError
from core.virtual_environment.__base.helpers import submit_command, with_connection
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
from core.workflow.node_events import NodeRunResult
@ -80,42 +80,41 @@ class CommandNode(Node[CommandNodeData]):
)
timeout = COMMAND_NODE_TIMEOUT_SECONDS if COMMAND_NODE_TIMEOUT_SECONDS > 0 else None
connection_handle = sandbox.establish_connection()
try:
command = shlex.split(raw_command)
with with_connection(sandbox) as conn:
command = shlex.split(raw_command)
sandbox_debug("command_node", "command", command)
sandbox_debug("command_node", "command", command)
future = sandbox.run_command(connection_handle, command, cwd=working_directory)
result = future.result(timeout=timeout)
future = submit_command(sandbox, conn, command, cwd=working_directory)
result = future.result(timeout=timeout)
outputs: dict[str, Any] = {
"stdout": result.stdout.decode("utf-8", errors="replace"),
"stderr": result.stderr.decode("utf-8", errors="replace"),
"exit_code": result.exit_code,
"pid": result.pid,
}
process_data = {"command": command, "working_directory": working_directory}
outputs: dict[str, Any] = {
"stdout": result.stdout.decode("utf-8", errors="replace"),
"stderr": result.stderr.decode("utf-8", errors="replace"),
"exit_code": result.exit_code,
"pid": result.pid,
}
process_data = {"command": command, "working_directory": working_directory}
if result.exit_code not in (None, 0):
stderr_text = result.stderr.decode("utf-8", errors="replace")
error_message = f"{stderr_text}\n\nCommand exited with code {result.exit_code}"
return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
outputs=outputs,
process_data=process_data,
error=error_message,
error_type=CommandExecutionError.__name__,
)
if result.exit_code not in (None, 0):
error_message = (
f"{result.stderr.decode('utf-8', errors='replace')}\n\nCommand exited with code {result.exit_code}"
)
return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
status=WorkflowNodeExecutionStatus.SUCCEEDED,
outputs=outputs,
process_data=process_data,
error=error_message,
error_type=CommandExecutionError.__name__,
)
return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED,
outputs=outputs,
process_data=process_data,
)
except CommandTimeoutError:
return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
@ -135,9 +134,6 @@ class CommandNode(Node[CommandNodeData]):
error=str(e),
error_type=type(e).__name__,
)
finally:
with contextlib.suppress(Exception):
sandbox.release_connection(connection_handle)
@classmethod
def _extract_variable_selector_to_variable_mapping(