chore: import upstream snapshot with attribution

examples/claude_code/swebench_utils/docker_runtime.py

@@ -0,0 +1,430 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Docker runtime management for repository setup and command execution.
+
+Provides containerized environment for repository testing with command execution,
+file operations, and state management capabilities.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import queue
+import re
+import threading
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+from docker.errors import DockerException, ImageNotFound
+from docker.models.containers import Container
+from swebench.harness.constants import SWEbenchInstance
+from typing_extensions import Self
+
+import docker
+
+# This will log to the console for debugging purposes.
+claude_code_logger = logging.getLogger("claude_code_agent.docker_runtime")
+
+CMD_OUTPUT_PS1_BEGIN = "\n###PS1JSON###\n"
+CMD_OUTPUT_PS1_END = "\n###PS1END###"
+CMD_OUTPUT_METADATA_PS1_REGEX = re.compile(
+    r"(?m)^\s*" + re.escape(CMD_OUTPUT_PS1_BEGIN.strip()) + r"\s*(.*?)\s*" + re.escape(CMD_OUTPUT_PS1_END.strip()),
+    re.DOTALL,
+)
+ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
+TIMEOUT_EXIT_CODE = 124
+
+MEM_LIMIT = "8g"
+CPU_CORES = 4
+
+
+VAR_PATTERNS = {
+    "exit_code": re.compile(r'"exit_code":\s*(-?\d+)\s*(?:,|\})'),
+    "username": re.compile(r'"username":\s*"([^"]*)"'),
+    "hostname": re.compile(r'"hostname":\s*"([^"]*)"'),
+    "working_dir": re.compile(r'"working_dir":\s*"([^"]*)"'),
+    "py_interpreter_path": re.compile(r'"py_interpreter_path":\s*"([^"]*)"'),
+}
+
+
+@dataclass
+class CmdOutputMetadata:
+    """
+    Additional metadata captured from PS1 shell prompt.
+
+    Provides context about command execution environment including
+    exit codes, user info, working directory, and Python interpreter.
+    """
+
+    exit_code: int = -1
+    username: str | None = None
+    hostname: str | None = None
+    working_dir: str | None = None
+    py_interpreter_path: str | None = None
+
+    @classmethod
+    def matches_ps1_metadata(cls, output: str) -> List[re.Match[str]]:
+        matches: List[re.Match[str]] = []
+        for match in CMD_OUTPUT_METADATA_PS1_REGEX.finditer(output):
+            scope = match.group(1).strip()
+            try:
+                d = json.loads(scope)  # Try to parse as JSON
+                matches.append(match)
+            except json.JSONDecodeError:
+                d = cls.best_effort_match(scope)
+                if len(d) > 0:
+                    matches.append(match)
+        return matches
+
+    @classmethod
+    def best_effort_match(cls, scope: str) -> Dict[str, Any]:
+        out: Dict[str, str] = {}
+        for field, pattern in VAR_PATTERNS.items():
+            m = pattern.search(scope)
+            if m:
+                out[field] = m.group(1)
+            else:
+                out[field] = ""
+        return out
+
+    @classmethod
+    def from_ps1_match(cls, match: re.Match[str]) -> Self:
+        """
+        Extract metadata from a PS1 prompt regex match.
+
+        Args:
+            match (re.Match[str]): Regex match containing JSON metadata
+
+        Returns:
+            Self: CmdOutputMetadata instance with parsed values
+        """
+        try:
+            metadata = json.loads(match.group(1))
+        except:
+            metadata = cls.best_effort_match(match.group(1))
+        # Create a copy of metadata to avoid modifying the original
+        processed = metadata.copy()
+        # Convert numeric fields
+        if "exit_code" in metadata:
+            try:
+                processed["exit_code"] = int(float(str(metadata["exit_code"])))
+            except (ValueError, TypeError):
+                processed["exit_code"] = -1
+        return cls(**processed)
+
+
+@dataclass
+class CommandResult:
+    """
+    Result of a command execution with output and metadata.
+
+    Attributes:
+        output (str): Command output text
+        metadata (Optional[CmdOutputMetadata]): Execution context metadata
+    """
+
+    output: str
+    metadata: Optional[CmdOutputMetadata]
+
+    def to_observation(self, strip: bool = True) -> str:
+        """
+        Convert command result to formatted observation string.
+
+        Args:
+            strip (bool): Whether to truncate long output
+
+        Returns:
+            str: Formatted observation with output and context
+        """
+        # compile regex once for efficiency
+        ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
+        output = ANSI_ESCAPE.sub("", self.output).replace("\r", "")
+
+        if len(output) > 1024 * 8 and strip:
+            output = output[: 1024 * 4] + "....stripped due to length....\n" + output[-1024 * 4 :]
+
+        if self.metadata is None:
+            return f"\n{output}\n"
+        return f"""{output}
+{self.metadata.username}@{self.metadata.hostname}:{self.metadata.working_dir} $
+
+exit code: {self.metadata.exit_code}
+"""
+
+
+class Runtime:
+    """
+    Docker container runtime for repository setup and testing.
+
+    Manages a Docker container with persistent bash session, command execution,
+    file operations, and container lifecycle management.
+    """
+
+    def __init__(self, container: Container, log_function: Callable[..., None]) -> None:
+        """
+        Initialize runtime with an existing Docker container.
+
+        Args:
+            container (Container): Docker container instance to manage
+        """
+        self.container = container
+        self.logger = log_function  # Set logger early so it's available even if init fails later
+        self.sock: Any = self.container.attach_socket(params={"stdin": 1, "stdout": 1, "stderr": 1, "stream": 1})  # type: ignore
+        self.output_queue: queue.Queue[bytes] = queue.Queue()
+        self._start_output_thread()
+        self._clear_initial_prompt()
+
+        json_str = json.dumps(
+            {
+                "exit_code": "$?",
+                "username": r"\u",
+                "hostname": r"\h",
+                "working_dir": r"$(pwd)",
+                "py_interpreter_path": r'$(which python 2>/dev/null || echo "")',
+            },
+            indent=2,
+        ).replace('"', r"\"")
+        ps1 = CMD_OUTPUT_PS1_BEGIN + json_str + CMD_OUTPUT_PS1_END + "\n"
+        self.send_command(f'export PROMPT_COMMAND=\'export PS1="{ps1}"\'; export PS2=""')
+        self.send_command("apt update -qq && apt install -y -qq git")
+        self.stopped = False
+
+    def _stream_output(self):
+        while True:
+            try:
+                output = self._recv_bytes(4096)
+                if not output:
+                    break
+                self.output_queue.put(output)
+            except (OSError, ConnectionError) as e:
+                print(f"Connection error in _stream_output: {e}")
+                break
+            except Exception as e:
+                # print(f"Unexpected error in _stream_output: {e}")
+                break
+
+    def _start_output_thread(self):
+        self.output_thread = threading.Thread(target=self._stream_output, daemon=True)
+        self.output_thread.start()
+        # TODO: kill the thread if main thread is stopped
+
+    def _clear_initial_prompt(self):
+        time.sleep(0.5)
+        while not self.output_queue.empty():
+            self.output_queue.get()
+
+    def _read_raw_output(self, timeout: float = 30) -> tuple[str, Optional[CmdOutputMetadata]]:
+        accumulated_output = ""
+        start_time = time.time()
+
+        while time.time() - start_time < timeout:
+            try:
+                chunk = self.output_queue.get(timeout=0.1)
+                accumulated_output += chunk.decode("utf-8", errors="ignore")
+                # PSReadLine injects ANSI + cursor control; normalize before matching
+                accumulated_clean = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
+                ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_clean)
+                if ps1_matches:
+                    break
+            except queue.Empty:
+                continue
+        accumulated_output = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
+        ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_output)
+        metadata = CmdOutputMetadata.from_ps1_match(ps1_matches[-1]) if ps1_matches else None
+        output = self._combine_outputs_between_matches(
+            accumulated_output,
+            ps1_matches,
+        )
+        return output, metadata
+
+    def _combine_outputs_between_matches(self, pane_content: str, ps1_matches: list[re.Match[str]]) -> str:
+        if len(ps1_matches) == 1:
+            return pane_content[: ps1_matches[0].start()]
+        elif len(ps1_matches) == 0:
+            return pane_content
+        output_segments: List[str] = []
+        for i in range(len(ps1_matches) - 1):
+            output_segment = pane_content[ps1_matches[i].end() + 1 : ps1_matches[i + 1].start()]
+            output_segments.append(output_segment)
+        return "\n".join(output_segments) + "\n" if output_segments else ""
+
+    def _recv_bytes(self, n: int = 4096) -> bytes:
+        # Prefer the public API on whatever object the SDK returns
+        for m in ("recv", "read"):
+            if hasattr(self.sock, m):
+                return getattr(self.sock, m)(n)
+        # Last-resort fallback for odd wrappers that still expose ._sock
+        if hasattr(self.sock, "_sock"):
+            for m in ("recv", "read"):
+                if hasattr(self.sock._sock, m):
+                    return getattr(self.sock._sock, m)(n)
+        raise TypeError(f"Don't know how to read from {type(self.sock).__name__}")
+
+    def _send_bytes(self, data: bytes) -> None:
+        if hasattr(self.sock, "_sock"):
+            for m in ("send", "sendall", "write"):
+                if hasattr(self.sock._sock, m):
+                    getattr(self.sock._sock, m)(data)
+                    return
+        for m in ("send", "sendall", "write"):
+            if hasattr(self.sock, m):
+                getattr(self.sock, m)(data)
+                return
+
+        raise TypeError(f"Don't know how to write to {type(self.sock).__name__}")
+
+    def _log_command_result(self, result: CommandResult) -> None:
+        claude_code_logger.debug("Docker runtime command finished with metadata: %s", result.metadata)
+        if len(result.output) > 2048:
+            logged_output = result.output[:1024] + "\n(... stripped due to length ...)\n" + result.output[-1024:]
+        else:
+            logged_output = result.output
+        claude_code_logger.debug(
+            "Docker runtime command finished with output (length = %d):\n%s", len(result.output), logged_output
+        )
+        # Output to the evaluation logger simultaneously
+        self.logger(text=logged_output)
+
+    def send_command(self, command: str, timeout: float = 20 * 60) -> CommandResult:
+        # Redact sensitive API keys from the command before logging
+        redacted_command = command
+        for sensitive_var in ["ANTHROPIC_AUTH_TOKEN", "API_KEY", "SECRET_KEY"]:
+            pattern = rf"(export (.*?){re.escape(sensitive_var)}(.*?)=)[^\s]+"
+            redacted_command = re.sub(pattern, rf"\1****REDACTED****", redacted_command)
+        claude_code_logger.info("Docker runtime receiving command: %s", redacted_command)
+        # Normalize newline semantics for interactive shells
+        if not command.endswith("\n"):
+            command += "\n"
+
+        while not self.output_queue.empty():
+            self.output_queue.get()
+
+        self._send_bytes(command.encode())
+
+        output, metadata = self._read_raw_output(timeout=timeout)
+        # TODO: Check exit code of the command (claude code download fail will not be caught by this)
+        if metadata is not None:
+            result = CommandResult(output=output, metadata=metadata)
+            self._log_command_result(result)
+            return result
+
+        # handle timeout
+        self._send_bytes(b"\x03")
+
+        kill_timeout = 5.0
+        kill_output, kill_metadata = self._read_raw_output(timeout=kill_timeout)
+
+        output = output + kill_output + "\n**Exited due to timeout**\n"
+        if kill_metadata is not None:
+            kill_metadata.exit_code = TIMEOUT_EXIT_CODE
+            result = CommandResult(output=output, metadata=kill_metadata)
+            self._log_command_result(result)
+            return result
+
+        fallback_metadata = CmdOutputMetadata(
+            exit_code=TIMEOUT_EXIT_CODE,
+        )
+        result = CommandResult(output=output, metadata=fallback_metadata)
+        self._log_command_result(result)
+        return result
+
+    def cleanup(self) -> None:
+        if self.stopped:
+            return
+        try:
+            claude_code_logger.info(f"Stopping container: {self.container.id}")
+            self.container.stop()
+            claude_code_logger.info(f"Removing container: {self.container.id}")
+            self.container.remove(force=True)
+            claude_code_logger.info(f"Container removed: {self.container.id}")
+            self.stopped = True
+        except Exception as e:
+            print(f"Failed to stop container: {e}")
+
+    def __del__(self):
+        self.cleanup()
+
+    @staticmethod
+    def pull_image(image_name: str) -> bool:
+        """
+        Pull Docker image from registry.
+
+        Args:
+            image_name (str): Name of the Docker image to pull
+
+        Returns:
+            bool: True if successful, False if image not found
+        """
+        client = docker.from_env()
+        try:
+            client.images.pull(image_name)
+            return True
+        except ImageNotFound:
+            return False
+
+    @classmethod
+    def start_session(
+        cls,
+        image_name: str,
+        instance: SWEbenchInstance,
+        log_function: Callable[..., None] = lambda: None,
+    ) -> Runtime:
+        """
+        Start a Docker container session for repository testing.
+
+        Args:
+            image_name (str): Base Docker image name
+            instance (dict): SWE-bench instance data with repo info
+
+        Returns:
+            SetupRuntime: Configured runtime session ready for command execution
+
+        Raises:
+            RuntimeError: If Docker is not available
+        """
+        try:
+            docker.from_env().ping()  # type: ignore
+        except DockerException:
+            raise RuntimeError("Docker is not installed or not running.")
+
+        _ = cls.pull_image(image_name)
+        client = docker.from_env(timeout=600)
+        container_id = instance["instance_id"]
+        container_name = f"git-launch-{container_id}-{str(uuid.uuid4())[:4]}"
+        info: Dict[str, str] = client.version()  # type: ignore
+        engine_os: str = (info.get("Os") or info.get("OSType") or "").lower()  # type: ignore
+        # which operating system this code is running on, note windows can run linux containers, so engine_os != (container) platform
+        extra_hosts = {"host.docker.internal": "host-gateway"} if "linux" in engine_os else None
+
+        shell_command = "/bin/bash"
+        working_dir = "/testbed"
+
+        claude_code_logger.info(
+            f"Starting container {container_name} with image {image_name}. Shell command: {shell_command}"
+        )
+        container = client.containers.run(
+            image_name,
+            name=container_name,
+            command=shell_command,
+            stdin_open=True,
+            tty=True,
+            detach=True,
+            environment={
+                "TERM": "xterm-mono",
+            },
+            working_dir=working_dir,
+            extra_hosts=extra_hosts,
+            network_mode="host",
+            cpu_quota=int(CPU_CORES * 100000),
+            mem_limit=MEM_LIMIT,
+        )
+        claude_code_logger.info(f"Container {container_name} started with ID: {container.id}")
+
+        session = cls(container, log_function=log_function)
+
+        return session