wehub-resource-sync
243 lines
11 KiB
Python
243 lines
11 KiB
Python
import os
|
|
import subprocess
|
|
import uuid
|
|
|
|
import docker
|
|
import pexpect
|
|
from docker.errors import DockerException, ImageNotFound, NotFound
|
|
|
|
|
|
class DockerManager:
|
|
"""
|
|
Manages Docker container lifecycle and command execution for the agent.
|
|
Supports both stateless (non-interactive) and stateful (interactive) modes.
|
|
"""
|
|
|
|
CONTAINER_TOOLS_PATH = "/agent_tools"
|
|
|
|
def __init__(
|
|
self,
|
|
image: str | None,
|
|
container_id: str | None,
|
|
dockerfile_path: str | None,
|
|
docker_image_file: str | None,
|
|
workspace_dir: str | None = None,
|
|
tools_dir: str | None = None,
|
|
interactive: bool = False,
|
|
):
|
|
if not image and not container_id and not dockerfile_path and not docker_image_file:
|
|
raise ValueError(
|
|
"Either a Docker image or a container ID or a dockerfile path or a docker image file (tar) must be provided."
|
|
)
|
|
self.client = docker.from_env()
|
|
self.image = image
|
|
self.container_id = container_id
|
|
self.dockerfile_path = dockerfile_path
|
|
self.docker_image_file = docker_image_file
|
|
self.workspace_dir = workspace_dir
|
|
self.tools_dir = tools_dir
|
|
self.interactive = interactive
|
|
self.container_workspace = "/workspace"
|
|
self.container = None
|
|
self.shell = None
|
|
self._is_managed = True
|
|
|
|
def start(self):
|
|
"""Starts/attaches to the container, mounts the workspace, copies tools, and starts the shell."""
|
|
try:
|
|
if self.dockerfile_path:
|
|
if not os.path.isabs(self.dockerfile_path):
|
|
raise ValueError("Dockerfile path must be an absolute path.")
|
|
build_context = os.path.dirname(self.dockerfile_path)
|
|
dockerfile_name = os.path.basename(self.dockerfile_path)
|
|
unique_tag = f"trae-agent-custom:{uuid.uuid4()}"
|
|
print(
|
|
f"Building Docker image from '{self.dockerfile_path}' with tag '{unique_tag}'..."
|
|
)
|
|
try:
|
|
new_image, build_logs = self.client.images.build(
|
|
path=build_context, dockerfile=dockerfile_name, tag=unique_tag, rm=True
|
|
)
|
|
self.image = new_image.tags[0]
|
|
print(f"✅ Successfully built image: {self.image}")
|
|
except Exception as e:
|
|
print("[red]❌ Docker image build failed. See logs below:[/red]")
|
|
for log_line in e.build_log:
|
|
if "stream" in log_line:
|
|
print(log_line["stream"].strip())
|
|
raise
|
|
|
|
elif self.docker_image_file:
|
|
print(f"Loading Docker image from file '{self.docker_image_file}'...")
|
|
try:
|
|
with open(self.docker_image_file, "rb") as f:
|
|
loaded_images = self.client.images.load(f.read())
|
|
if not loaded_images:
|
|
raise DockerException("Failed to load any images from the provided file.")
|
|
self.image = loaded_images[0].tags[0]
|
|
print(f"✅ Successfully loaded image: {self.image}")
|
|
except FileNotFoundError:
|
|
raise
|
|
except Exception as e:
|
|
raise DockerException(f"Error loading image from file: {e}") from e
|
|
|
|
if self.container_id:
|
|
print(f"Attaching to existing container: {self.container_id}...")
|
|
self.container = self.client.containers.get(self.container_id)
|
|
self._is_managed = False
|
|
print(f"Successfully attached to container {self.container.short_id}.")
|
|
elif self.image:
|
|
print(f"Starting a new container from image: {self.image}...")
|
|
if self.workspace_dir is not None:
|
|
os.makedirs(self.workspace_dir, exist_ok=True)
|
|
volumes = {
|
|
os.path.abspath(self.workspace_dir): {
|
|
"bind": self.container_workspace,
|
|
"mode": "rw",
|
|
}
|
|
}
|
|
self.container = self.client.containers.run(
|
|
self.image,
|
|
command="sleep infinity",
|
|
detach=True,
|
|
volumes=volumes,
|
|
working_dir=self.container_workspace,
|
|
)
|
|
self.container_id = self.container.id
|
|
self._is_managed = True
|
|
print(
|
|
f"Container {self.container.short_id} created. Workspace '{self.workspace_dir}' is mounted to '{self.container_workspace}'."
|
|
)
|
|
else:
|
|
self.container = self.client.containers.run(
|
|
self.image,
|
|
command="sleep infinity",
|
|
detach=True,
|
|
working_dir=self.container_workspace,
|
|
)
|
|
self.container_id = self.container.id
|
|
self._is_managed = True
|
|
print(f"Container {self.container.short_id} created.")
|
|
self._copy_tools_to_container()
|
|
# if self.interactive:
|
|
self._start_persistent_shell()
|
|
except (ImageNotFound, NotFound, DockerException) as e:
|
|
print(f"[red]Failed to start DockerManager: {e}[/red]")
|
|
raise
|
|
|
|
def execute(self, command: str, timeout: int = 300) -> tuple[int, str]:
|
|
"""
|
|
Executes a command using the configured mode (interactive or stateless).
|
|
"""
|
|
if not self.container:
|
|
raise RuntimeError("Container is not running. Call start() first.")
|
|
|
|
# if self.interactive:
|
|
return self._execute_interactive(command, timeout)
|
|
# else:
|
|
# return self._execute_stateless(command)
|
|
|
|
def stop(self):
|
|
"""Stops the pexpect shell and cleans up the container if managed by this instance."""
|
|
if self.shell and self.shell.isalive():
|
|
print("Closing persistent shell...")
|
|
self.shell.close(force=True)
|
|
self.shell = None
|
|
|
|
if self.container and self._is_managed:
|
|
print(f"Stopping and removing managed container {self.container.short_id}...")
|
|
try:
|
|
self.container.stop()
|
|
self.container.remove()
|
|
print("Container cleaned up successfully.")
|
|
except DockerException as e:
|
|
print(
|
|
f"[yellow]Warning: Could not clean up container {self.container.short_id}: {e}[/yellow]"
|
|
)
|
|
|
|
self.container = None
|
|
|
|
# --- Private Helper Methods ---
|
|
|
|
def _copy_tools_to_container(self):
|
|
"""Copies the local tools directory to a fixed path inside the container."""
|
|
if not self.tools_dir or not os.path.isdir(self.tools_dir):
|
|
print(
|
|
f"[yellow]Packaged tools directory '{self.tools_dir}' not provided or not found, skipping copy.[/yellow]"
|
|
)
|
|
return
|
|
|
|
print(
|
|
f"Copying tools from '{self.tools_dir}' to container path '{self.CONTAINER_TOOLS_PATH}'..."
|
|
)
|
|
try:
|
|
cmd = f"docker cp '{os.path.abspath(self.tools_dir)}' '{self.container.id}:{self.CONTAINER_TOOLS_PATH}'"
|
|
subprocess.run(cmd, shell=True, check=True, capture_output=True)
|
|
print("Tools copied successfully.")
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"[red]Failed to copy tools to container: {e.stderr.decode()}[/red]")
|
|
raise DockerException(f"Failed to copy tools: {e.stderr.decode()}") from e
|
|
|
|
def _start_persistent_shell(self):
|
|
"""Spawns a persistent bash shell inside the container using pexpect."""
|
|
if not self.container:
|
|
return
|
|
# print("Starting persistent shell for interactive mode...")
|
|
try:
|
|
command = f"docker exec -it {self.container.id} /bin/bash"
|
|
self.shell = pexpect.spawn(command, encoding="utf-8", timeout=120)
|
|
self.shell.expect([r"\$", r"#"], timeout=120)
|
|
print("Persistent shell is ready.")
|
|
except pexpect.exceptions.TIMEOUT:
|
|
print(
|
|
"[red]Timeout waiting for shell prompt. The container might be slow to start or misconfigured.[/red]"
|
|
)
|
|
raise
|
|
|
|
# def _execute_stateless(self, command: str) -> tuple[int, str]:
|
|
# """Executes a command in a new, non-persistent session."""
|
|
# print(f"Executing (stateless): `{command}`")
|
|
# exit_code, output_bytes = self.container.exec_run(cmd=f"/bin/sh -c '{command}'")
|
|
# output = output_bytes.decode('utf-8', errors='replace').strip()
|
|
# return exit_code, output
|
|
|
|
def _execute_interactive(self, command: str, timeout: int) -> tuple[int, str]:
|
|
"""Executes a command within the existing persistent shell."""
|
|
if not self.shell or not self.shell.isalive():
|
|
print("[yellow]Shell not found or died. Attempting to restart...[/yellow]")
|
|
self._start_persistent_shell()
|
|
|
|
if self.shell is None:
|
|
raise RuntimeError("Failed to start or restart the persistent shell.")
|
|
|
|
marker = "---CMD_DONE---"
|
|
full_command = command.strip()
|
|
marker_command = f"echo {marker}$?"
|
|
self.shell.sendline(full_command)
|
|
self.shell.sendline(marker_command)
|
|
try:
|
|
self.shell.expect(marker + r"(\d+)", timeout=timeout)
|
|
except pexpect.exceptions.TIMEOUT:
|
|
return (
|
|
-1,
|
|
f"Error: Command '{command}' timed out after {timeout} seconds. Partial output:\n{self.shell.before}",
|
|
)
|
|
exit_code = int(self.shell.match.group(1))
|
|
|
|
output_before_marker = self.shell.before
|
|
|
|
# 1. Split the raw output into lines
|
|
all_lines = output_before_marker.splitlines()
|
|
# 2. Filter out the lines that are just echoes of our commands
|
|
clean_lines = []
|
|
for line in all_lines:
|
|
stripped_line = line.strip()
|
|
# Ignore the line if it's an echo of the original command OR our marker command
|
|
if stripped_line != full_command and marker_command not in stripped_line:
|
|
clean_lines.append(line)
|
|
# 3. Join the clean lines back together
|
|
cleaned_output = "\n".join(clean_lines)
|
|
# Wait for the next shell prompt to ensure the shell is ready
|
|
self.shell.expect([r"\$", r"#"])
|
|
return exit_code, cleaned_output.strip()
|