From 34382e6cc961d9f2354280f07325eacbb9279420 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Mon, 27 Jul 2026 12:08:42 -0700 Subject: [PATCH 01/13] feat(sandbox): add Amazon Bedrock AgentCore Runtime backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `--sandbox agentcore`, running each rollout in an AgentCore Runtime microVM, and first refactors the transport layer the new backend would otherwise have had to copy. Refactor (behavior-preserving) Transport selection was an `if environment == "..."` chain in the ACP layer that re-derived the backend from a string while already holding the sandbox object, and sniffed private attributes (`env._strategy`) to tell Daytona's variants apart. Backends now answer `BaseSandbox.live_process(agent=...)` themselves, so the chain collapses to one call and a new backend adds no branch. This also makes an existing latent bug explicit: Modal fell through the chain's `else` and was handed a `DaytonaProcess`, failing deep inside Daytona SSH setup; it now raises an actionable NotImplementedError. `sandbox/process.py` had reached 1026 lines and is split into a package (import paths unchanged). `LiveProcess` becomes a pure ABC and the local subprocess implementation moves to `SubprocessLiveProcess`. That split is what the two WebSocket transports need: `DaytonaPtyProcess` already had to neutralize the base class with `_process = None # Not used`, and AgentCore is the second such transport. Per-backend capability facts move onto the provider registry, so the no-network gate reads `enforces_no_network` instead of growing a `sandbox == ""` special case per backend that cannot isolate. Also drops a stale hand-copied provider list from a CLI docstring. AgentCore backend AgentCore takes no image per run — an image must be pushed to ECR and registered as an agent runtime. BenchFlow provisions per task: build for linux/arm64, push under a content-digest tag, reuse or create the runtime. The digest tag is what makes repeat runs cheap (measured: environment_setup 58.9s cold, 16.6s on reuse). Commands go through InvokeAgentRuntimeCommand; the ACP agent runs on the `open_shell` WebSocket, which shares a session and filesystem with exec. Three platform behaviors were established by experiment against the live service, not from docs, and are encoded here: - A container that does not answer `GET /ping` on :8080 reaches READY but 500s on every command. Task images know nothing about AgentCore, so a stdlib-only responder is appended to the image as its entrypoint. - READY describes the runtime definition, not a running session. The first command on a cold, never-pulled image 500s while the container boots, then succeeds. Warm-up retries with backoff so that race cannot surface as a scored-0 rollout. - The shell is a PTY (echo, bracketed paste, CRLF), the same shape as DaytonaPtyProcess, and is handled the same proven way. Gated rather than faked: snapshots (no primitive), multi-service compose (single container), and `network_mode = "no-network"` — `networkMode` offers only PUBLIC or VPC, so no-network tasks are refused instead of run unisolated. Throttling and quota errors map to SandboxStartupError so capacity problems stay attributable as infrastructure. --- docs/running-benchmarks.md | 37 + pyproject.toml | 9 + src/benchflow/acp/runtime.py | 37 +- src/benchflow/cli/main.py | 9 +- src/benchflow/sandbox/_base.py | 27 +- src/benchflow/sandbox/agentcore.py | 860 ++++++++++++++++++ src/benchflow/sandbox/apple_container.py | 9 + src/benchflow/sandbox/daytona.py | 22 +- src/benchflow/sandbox/docker.py | 10 +- src/benchflow/sandbox/process/__init__.py | 39 + src/benchflow/sandbox/process/_base.py | 193 ++++ src/benchflow/sandbox/process/agentcore.py | 260 ++++++ src/benchflow/sandbox/process/apple.py | 101 ++ .../{process.py => process/daytona.py} | 386 +------- src/benchflow/sandbox/process/docker.py | 157 ++++ src/benchflow/sandbox/providers.py | 21 + src/benchflow/sandbox/setup.py | 17 + src/benchflow/task/runtime_capabilities.py | 10 +- tests/test_acp.py | 264 +----- tests/test_acp_model_config_dispatch.py | 4 - tests/test_acp_setup_failure_propagation.py | 28 - tests/test_agentcore_sandbox.py | 419 +++++++++ tests/test_oracle_chokepoint.py | 17 +- tests/test_process.py | 4 +- tests/test_rollout_on_ask_user_wiring.py | 4 - tests/test_runtime_capabilities.py | 3 +- tests/test_sandbox_live_process.py | 224 +++++ tests/test_sandbox_provider_registry_drift.py | 20 +- uv.lock | 45 +- 29 files changed, 2519 insertions(+), 717 deletions(-) create mode 100644 src/benchflow/sandbox/agentcore.py create mode 100644 src/benchflow/sandbox/process/__init__.py create mode 100644 src/benchflow/sandbox/process/_base.py create mode 100644 src/benchflow/sandbox/process/agentcore.py create mode 100644 src/benchflow/sandbox/process/apple.py rename src/benchflow/sandbox/{process.py => process/daytona.py} (63%) create mode 100644 src/benchflow/sandbox/process/docker.py create mode 100644 tests/test_agentcore_sandbox.py create mode 100644 tests/test_sandbox_live_process.py diff --git a/docs/running-benchmarks.md b/docs/running-benchmarks.md index 433b5469c..1988506f3 100644 --- a/docs/running-benchmarks.md +++ b/docs/running-benchmarks.md @@ -347,6 +347,7 @@ The **Harvey LAB harness** agent is special — it runs Harvey LAB's own agent l | Apple Container | `--sandbox apple-container` | Local Apple Silicon macOS runs without Docker Desktop | | Daytona | `--sandbox daytona` | Cloud runs with concurrency (needs `DAYTONA_API_KEY`) | | Modal | `--sandbox modal` | Serverless, high concurrency (needs Modal auth) | +| AgentCore | `--sandbox agentcore` | AWS-native isolated microVMs (needs AWS credentials) | Apple Container requires Apple Container 1.1+ on Apple Silicon and runs the model proxy inside each VM. It supports public-network, single-container arm64 tasks and @@ -356,6 +357,42 @@ running concurrent BenchFlow processes, because the macOS allocation leak is system-wide. Use Docker, Daytona, or Modal for `network_mode = "no-network"`, multi-service, snapshot, or high-concurrency runs. +### Amazon Bedrock AgentCore + +`--sandbox agentcore` runs each rollout in an AgentCore Runtime microVM. +Install the extra and point BenchFlow at an execution role: + +```bash +pip install 'benchflow[sandbox-agentcore]' +export BENCHFLOW_AGENTCORE_ROLE_ARN="arn:aws:iam:::role/" +export BENCHFLOW_AGENTCORE_REGION="us-west-2" # optional, this is the default +``` + +The role must be assumable by `bedrock-agentcore.amazonaws.com` and able to +pull from ECR and write CloudWatch logs. Your own credentials additionally need +`bedrock-agentcore:*` and ECR push permissions. + +Unlike the other backends, AgentCore cannot take a container image directly: +BenchFlow builds the task image for `linux/arm64`, pushes it to ECR (repository +`benchflow-agentcore`, override with `BENCHFLOW_AGENTCORE_ECR_REPOSITORY`), and +registers it as an agent runtime. The image is tagged by content digest, so a +task whose environment has not changed reuses the pushed layers and the +existing runtime instead of paying for a rebuild. + +BenchFlow also appends a small HTTP responder to the image and makes it the +entrypoint. AgentCore refuses command execution for a session whose container +does not answer `GET /ping` on port 8080, and task images know nothing about +AgentCore, so this shim is what makes an ordinary task image runnable. + +Constraints: `linux/arm64` only, single container (no compose/multi-service +tasks), no snapshot support, and `network_mode = "no-network"` is **not** +enforceable — AgentCore's network mode is either `PUBLIC` or `VPC`, so +BenchFlow refuses no-network tasks on this backend rather than running them +unisolated. The model proxy runs inside the sandbox, as on Daytona and Modal. +Sessions default to a 15-minute idle timeout and an 8-hour lifetime; override +with `BENCHFLOW_AGENTCORE_IDLE_TIMEOUT_SEC` / `BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC` +if agent turns are long enough to risk reclamation mid-run. + For large-scale runs (100+ tasks), use Daytona or Modal with high concurrency: ```bash diff --git a/pyproject.toml b/pyproject.toml index c6a6fb256..898c316a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,15 @@ sandbox-modal = [ "modal>=0.73", "tenacity>=8.0", ] +sandbox-agentcore = [ + # Bedrock AgentCore Runtime microVMs. The SDK supplies the interactive + # shell (`open_shell`) used as the ACP transport; boto3 supplies the + # control plane (CreateAgentRuntime) and the command plane + # (InvokeAgentRuntimeCommand). The boto3 floor is the first release that + # models the bedrock-agentcore/-control services. + "bedrock-agentcore>=1.18", + "boto3>=1.43.31", +] bedrock = [ "boto3>=1.40", ] diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 1443b874d..0653dda41 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -46,12 +46,6 @@ build_priv_drop_cmd, enforce_agent_egress_firewall, ) -from benchflow.sandbox.process import ( - AppleContainerProcess, - DaytonaProcess, - DaytonaPtyProcess, - DockerProcess, -) from benchflow.trajectories._capture import _capture_session_trajectory # Re-exported for backwards compatibility — tests and downstream code @@ -577,30 +571,13 @@ async def connect_acp( await asyncio.sleep(delay) try: - if environment == "docker": - live_proc = DockerProcess.from_sandbox_env(env) - elif environment == "apple-container": - live_proc = AppleContainerProcess.from_sandbox_env(env) - elif environment == "daytona": - transport_name = selected_acp_transport( - agent=agent, - environment=environment, - ) - if transport_name == "ssh": - live_proc = await DaytonaProcess.from_sandbox_env(env) - logger.info("Using SSH transport for %s on Daytona", agent) - else: - live_proc = await DaytonaPtyProcess.from_sandbox_env(env) - logger.info("Using PTY transport for Daytona sandbox") - else: - is_dind = hasattr(env, "_strategy") and hasattr( - env._strategy, "_compose_cmd" - ) - if is_dind: - live_proc = await DaytonaPtyProcess.from_sandbox_env(env) - logger.info("Using PTY transport for DinD compose task") - else: - live_proc = await DaytonaProcess.from_sandbox_env(env) + # The sandbox owns its transport: it knows how to reach into + # itself, so there is no provider-name branch here. Backends with + # no live-process transport raise an actionable NotImplementedError + # from BaseSandbox.live_process rather than silently receiving + # another backend's transport (Modal previously fell through to + # DaytonaProcess and failed deep inside SSH setup). + live_proc = await env.live_process(agent=agent) agent_log = rollout_dir / "agent" / f"{agent.replace('-', '_')}.txt" transport = ContainerTransport( diff --git a/src/benchflow/cli/main.py b/src/benchflow/cli/main.py index 7347b9190..df737153a 100644 --- a/src/benchflow/cli/main.py +++ b/src/benchflow/cli/main.py @@ -586,10 +586,11 @@ def eval_run( typer.Option("--trials", help="Number of trials for --matrix"), ] = 1, ) -> None: - """Run an evaluation — single task or batch. - - Sandbox: docker, daytona, modal, or apple-container. - """ + # The supported --sandbox values are rendered from the provider registry + # into that option's own help text. This docstring used to hand-copy them + # ("docker, daytona, modal, or apple-container") and went stale the moment + # a backend was added, so it no longer repeats them. + """Run an evaluation — single task or batch.""" _apply_dotenv_to_process_env() request = EvalCreateRequest( diff --git a/src/benchflow/sandbox/_base.py b/src/benchflow/sandbox/_base.py index aa627013b..80824f987 100644 --- a/src/benchflow/sandbox/_base.py +++ b/src/benchflow/sandbox/_base.py @@ -14,7 +14,7 @@ import shlex from abc import ABC, abstractmethod from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 from pydantic import BaseModel @@ -24,6 +24,9 @@ from benchflow.task.env import resolve_env_vars from benchflow.task.paths import RolloutPaths +if TYPE_CHECKING: + from benchflow.sandbox.process import LiveProcess + logger = logging.getLogger("benchflow") # Docker Compose service names are restricted to this grammar. Used to filter @@ -386,6 +389,28 @@ async def is_file( async def attach(self) -> None: raise NotImplementedError("This environment does not support attaching.") + async def live_process(self, *, agent: str | None = None) -> LiveProcess: + """Open the bidirectional line pipe an ACP agent speaks over. + + ``agent`` names the agent being launched; backends that pick a + transport per agent (Daytona: SSH for gemini, PTY otherwise) use it. + + Each backend knows how to reach into its own sandbox, so the transport + choice belongs here rather than in a provider-name ``if/elif`` at the + ACP layer — that chain had to re-derive the backend from a string while + already holding this object, and sniff private attributes + (``env._strategy``) to tell Daytona's variants apart. + + Backends with no live-process transport (Modal) leave this default in + place and fail with an actionable message instead of being handed + another backend's transport. + """ + raise NotImplementedError( + f"{type(self).__name__} does not provide a live agent transport, so " + "it cannot host an ACP agent. Use a backend that implements " + "live_process() (docker, daytona, apple-container, agentcore)." + ) + # Container-level snapshot/restore (Branch substrate) # # Part of the Sandbox contract (``docs/architecture.md``): the Branch diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py new file mode 100644 index 000000000..77bb6cf54 --- /dev/null +++ b/src/benchflow/sandbox/agentcore.py @@ -0,0 +1,860 @@ +"""Amazon Bedrock AgentCore Runtime sandbox backend. + +AgentCore runs each session in an isolated Firecracker-class microVM. Unlike +Docker/Modal/Daytona it does not take an image per run: an image must first be +pushed to ECR and registered as an *agent runtime* (a control-plane resource +with its own ARN), after which sessions are invoked against that ARN. + +BenchFlow therefore provisions per task: build the task's image for +``linux/arm64``, push it to ECR under a tag derived from the image's content +digest, and reuse (or create) the agent runtime bound to that tag. The digest +tag is what makes repeat runs cheap — the second task with an unchanged +environment reuses both the ECR layer set and the runtime. + +Two facts about the platform shape this module, both established by direct +experiment against the live service rather than from the docs: + +1. **The container must answer the Runtime HTTP contract.** A microVM whose + image merely sleeps reaches ``READY`` but every + ``InvokeAgentRuntimeCommand`` against it fails with a 500 from the runtime. + Serving ``GET /ping`` on port 8080 fixes it. Task images know nothing about + AgentCore, so BenchFlow appends a tiny stdlib-only responder to the image + and makes it the entrypoint (see ``_PING_SHIM``). +2. **Command execution and the interactive shell share one session.** A file + written by ``exec()`` is visible to the agent running under + ``open_shell()`` when both use the same ``runtimeSessionId``, which is what + lets the kernel stage files and the verifier read results around an agent + that lives on the WebSocket. + +Not supported, and gated rather than faked: container snapshots (no platform +primitive — the ``BaseSandbox`` default raises), multi-service compose +topologies (single container), and ``network_mode = "no-network"`` — the +``networkConfiguration.networkMode`` enum offers only ``PUBLIC`` and ``VPC``, +so isolation is declared unsupported in the provider registry instead of being +silently ignored. +""" + +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import io +import os +import re +import shlex +import shutil +import subprocess +import tarfile +import threading +import uuid +from pathlib import Path, PurePosixPath +from typing import TYPE_CHECKING, Any + +from benchflow._paths import iter_safe_tree +from benchflow.sandbox._base import BaseSandbox, ExecResult, wrap_command_with_env_file +from benchflow.sandbox.protocol import SandboxStartupError +from benchflow.task.config import SandboxConfig +from benchflow.task.paths import RolloutPaths + +if TYPE_CHECKING: + from benchflow.sandbox.process import LiveProcess + +_DEFAULT_REGION = "us-west-2" +_DEFAULT_ECR_REPOSITORY = "benchflow-agentcore" +_RUNTIME_READY_TIMEOUT_SEC = 600 +_RUNTIME_POLL_INTERVAL_SEC = 5 +# A cold image can take a while to pull into the microVM on the first command. +_SESSION_WARMUP_ATTEMPTS = 8 +_SESSION_WARMUP_BACKOFF_SEC = 4.0 +_SESSION_WARMUP_MAX_BACKOFF_SEC = 30.0 +# The service caps a single command payload at 64 KB. Base64 inflates by 4/3, +# and the wrapper adds shell scaffolding, so keep a wide margin. +_MAX_INLINE_UPLOAD_BYTES = 24 * 1024 +_MAX_DOWNLOAD_BYTES = 64 * 1024 * 1024 +_ECR_TAG_RE = re.compile(r"[^a-zA-Z0-9_.-]+") +# ``_PING_SHIM`` implements the plain HTTP contract (GET /ping, +# POST /invocations), so register the runtime as HTTP rather than relying on +# whatever the service currently defaults to. +_PROTOCOL_CONFIGURATION = {"serverProtocol": "HTTP"} + +# Environment overrides. +_ENV_REGION = "BENCHFLOW_AGENTCORE_REGION" +_ENV_ROLE_ARN = "BENCHFLOW_AGENTCORE_ROLE_ARN" +_ENV_ECR_REPOSITORY = "BENCHFLOW_AGENTCORE_ECR_REPOSITORY" +_ENV_IDLE_TIMEOUT = "BENCHFLOW_AGENTCORE_IDLE_TIMEOUT_SEC" +_ENV_MAX_LIFETIME = "BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC" + +# Stdlib-only responder for the AgentCore Runtime HTTP contract. Kept +# dependency-free so it runs on any task base image that has a `python3`. +_PING_SHIM = '''\ +"""BenchFlow shim: satisfies the AgentCore Runtime HTTP contract. + +AgentCore refuses to service InvokeAgentRuntimeCommand for a session whose +container does not answer this contract, so BenchFlow injects this responder +as the image entrypoint. It does no work beyond staying alive and replying; +the agent itself is launched later over the shell WebSocket. +""" +import json +from http.server import BaseHTTPRequestHandler, HTTPServer + + +class _Handler(BaseHTTPRequestHandler): + def _reply(self, code, payload): + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path.rstrip("/") in ("/ping", ""): + self._reply(200, {"status": "Healthy"}) + else: + self._reply(404, {"error": "not found"}) + + def do_POST(self): + self.rfile.read(int(self.headers.get("Content-Length") or 0)) + self._reply(200, {"result": "benchflow-sandbox"}) + + def log_message(self, *_args): + return + + +if __name__ == "__main__": + HTTPServer(("0.0.0.0", 8080), _Handler).serve_forever() +''' + + +def _run(*args: str, timeout: float | None = None, cwd: str | None = None): + return subprocess.run( + args, capture_output=True, text=True, timeout=timeout, cwd=cwd + ) + + +class AgentCoreSandbox(BaseSandbox): + """Sandbox backend for Bedrock AgentCore Runtime microVMs.""" + + def __init__( + self, + environment_dir: Path, + environment_name: str, + session_id: str, + rollout_paths: RolloutPaths | None, + task_env_config: SandboxConfig, + *args: Any, + **kwargs: Any, + ) -> None: + super().__init__( + environment_dir=environment_dir, + environment_name=environment_name, + session_id=session_id, + rollout_paths=rollout_paths, + task_env_config=task_env_config, + **kwargs, + ) + self.region = os.environ.get(_ENV_REGION) or _DEFAULT_REGION + self._ecr_repository = ( + os.environ.get(_ENV_ECR_REPOSITORY) or _DEFAULT_ECR_REPOSITORY + ) + self._role_arn = os.environ.get(_ENV_ROLE_ARN) + self.runtime_arn: str | None = None + self.runtime_session_id: str | None = None + self._runtime_id: str | None = None + self._created_runtime = False + self._account_id: str | None = None + self._clients: dict[str, Any] = {} + self._client_lock = threading.Lock() + + # ---------------------------------------------------------------- config + + @property + def is_mounted(self) -> bool: + return False + + @property + def sandbox_id(self) -> str | None: + return self.runtime_arn + + @classmethod + def preflight(cls) -> None: + try: + import bedrock_agentcore # noqa: F401 + import boto3 + except ImportError as exc: + raise SystemExit( + "AgentCore requires the 'sandbox-agentcore' extra. Install it " + "with `uv sync --extra sandbox-agentcore` or " + "`pip install 'benchflow[sandbox-agentcore]'`." + ) from exc + + from botocore.exceptions import BotoCoreError, ClientError + + try: + boto3.Session().client( + "sts", region_name=os.environ.get(_ENV_REGION) or _DEFAULT_REGION + ).get_caller_identity() + except (BotoCoreError, ClientError) as exc: + raise SystemExit( + "AgentCore requires working AWS credentials. Configure them with " + "`aws configure`, an SSO profile, or AWS_* environment " + f"variables. Underlying error: {exc}" + ) from exc + + if not shutil.which("docker"): + raise SystemExit( + "AgentCore builds task images locally and pushes them to ECR, " + "which requires the docker CLI on PATH." + ) + + def _validate_definition(self) -> None: + dockerfile = self.environment_dir / "Dockerfile" + if not dockerfile.exists() and not self.task_env_config.docker_image: + raise ValueError( + f"No Dockerfile found in {self.environment_dir} and no " + "docker_image specified in task config." + ) + + # ------------------------------------------------------------- lifecycle + + def _client(self, service: str) -> Any: + """Return a cached boto3 client for *service*. + + Building a client parses the service's JSON model and re-resolves + credentials, which is far too expensive to repeat per ``exec()`` — and + it floods the run log with botocore's credential-discovery line. The + cache is guarded because ``exec`` dispatches through + ``asyncio.to_thread``, so concurrent calls can arrive on different + threads. botocore clients are safe to share for API calls. + """ + cached = self._clients.get(service) + if cached is not None: + return cached + with self._client_lock: + cached = self._clients.get(service) + if cached is None: + import boto3 + + cached = boto3.Session(region_name=self.region).client(service) + self._clients[service] = cached + return cached + + def _resolve_account_id(self) -> str: + if self._account_id is None: + self._account_id = self._client("sts").get_caller_identity()["Account"] + return self._account_id + + def _ecr_registry(self) -> str: + return f"{self._resolve_account_id()}.dkr.ecr.{self.region}.amazonaws.com" + + def _image_tag(self) -> str: + safe = _ECR_TAG_RE.sub("_", self.environment_name)[:80] + return f"bf-{safe}".lower() + + async def start(self, force_build: bool) -> None: + try: + image_uri = await self._publish_image(force_build=force_build) + self.runtime_arn = await self._ensure_runtime(image_uri) + # AgentCore requires a session id of at least 33 characters. + self.runtime_session_id = f"{uuid.uuid4()}-{uuid.uuid4().hex[:8]}" + await self._warm_session() + self.logger.info( + "AgentCore session ready (runtime=%s, session=%s)", + self.runtime_arn, + self.runtime_session_id, + ) + except SandboxStartupError: + await self._safe_teardown() + raise + except Exception as exc: + await self._safe_teardown() + raise SandboxStartupError( + f"AgentCore sandbox failed to start: {exc}", + sandbox_id=self.runtime_arn, + build_timeout_sec=self.task_env_config.build_timeout_sec, + ) from exc + + async def _publish_image(self, *, force_build: bool) -> str: + """Build the task image with the ping shim and push it to ECR.""" + build_context, dockerfile, shim = self._materialize_build_context() + registry = self._ecr_registry() + repo = self._ecr_repository + local_tag = self._image_tag() + + build_args = [ + "docker", + "build", + "--platform", + "linux/arm64", + "-f", + str(dockerfile), + "-t", + local_tag, + ] + if force_build: + build_args.append("--no-cache") + build_args.append(str(build_context)) + + try: + result = await asyncio.to_thread( + _run, *build_args, timeout=self.task_env_config.build_timeout_sec + ) + finally: + # The generated Dockerfile and shim only exist for the build. The + # build context is the task's own environment directory, so leaving + # them behind would litter the caller's task tree. + for path in (dockerfile, shim): + path.unlink(missing_ok=True) + if result.returncode != 0: + raise RuntimeError( + f"docker build failed (exit {result.returncode}):\n" + f"{(result.stderr or result.stdout or '').strip()[:4000]}" + ) + + # Tag by image content digest so an unchanged environment reuses the + # already-pushed layers and the already-registered agent runtime. + digest = await asyncio.to_thread( + _run, "docker", "images", "--no-trunc", "--quiet", local_tag + ) + content_id = (digest.stdout or "").strip().replace("sha256:", "")[:32] + if not content_id: + content_id = hashlib.sha256(local_tag.encode()).hexdigest()[:32] + remote_tag = f"{local_tag}-{content_id}" + image_uri = f"{registry}/{repo}:{remote_tag}" + + await asyncio.to_thread(self._ensure_ecr_repository) + if not force_build and await asyncio.to_thread( + self._ecr_image_exists, remote_tag + ): + self.logger.info("Reusing published AgentCore image %s", image_uri) + return image_uri + + await asyncio.to_thread(self._docker_login, registry) + retag = await asyncio.to_thread(_run, "docker", "tag", local_tag, image_uri) + if retag.returncode != 0: + raise RuntimeError(f"docker tag failed: {(retag.stderr or '')[:2000]}") + push = await asyncio.to_thread(_run, "docker", "push", image_uri, timeout=1800) + if push.returncode != 0: + raise RuntimeError( + f"docker push to ECR failed (exit {push.returncode}):\n" + f"{(push.stderr or push.stdout or '').strip()[:4000]}" + ) + self.logger.info("Pushed AgentCore image %s", image_uri) + return image_uri + + def _materialize_build_context(self) -> tuple[Path, Path, Path]: + """Return (context_dir, dockerfile, shim) for an image with the shim. + + The task's own Dockerfile is left untouched on disk; the shim is + appended into a *generated* Dockerfile inside the same build context so + relative ``COPY`` paths keep resolving. Both generated files are + transient — the caller removes them once the build finishes. + """ + context = self.environment_dir + base_dockerfile = context / "Dockerfile" + if self.task_env_config.docker_image: + base = f"FROM {self.task_env_config.docker_image}\n" + else: + base = base_dockerfile.read_text() + + shim_path = context / ".benchflow_agentcore_shim.py" + shim_path.write_text(_PING_SHIM) + generated = context / "Dockerfile.benchflow-agentcore" + generated.write_text( + base + + "\n" + + "# --- BenchFlow AgentCore runtime contract ---\n" + + "# AgentCore refuses command execution for a session whose\n" + + "# container does not answer GET /ping on :8080.\n" + + "COPY .benchflow_agentcore_shim.py /opt/benchflow_agentcore_shim.py\n" + + "EXPOSE 8080\n" + + "ENTRYPOINT []\n" + + 'CMD ["python3", "/opt/benchflow_agentcore_shim.py"]\n' + ) + return context, generated, shim_path + + def _ensure_ecr_repository(self) -> None: + from botocore.exceptions import ClientError + + ecr = self._client("ecr") + try: + ecr.create_repository(repositoryName=self._ecr_repository) + except ClientError as exc: + if exc.response["Error"]["Code"] != "RepositoryAlreadyExistsException": + raise + + def _ecr_image_exists(self, tag: str) -> bool: + from botocore.exceptions import ClientError + + try: + self._client("ecr").describe_images( + repositoryName=self._ecr_repository, imageIds=[{"imageTag": tag}] + ) + return True + except ClientError: + return False + + def _docker_login(self, registry: str) -> None: + token = self._client("ecr").get_authorization_token() + blob = token["authorizationData"][0]["authorizationToken"] + _user, password = base64.b64decode(blob).decode().split(":", 1) + proc = subprocess.run( + ["docker", "login", "--username", "AWS", "--password-stdin", registry], + input=password, + capture_output=True, + text=True, + timeout=120, + ) + if proc.returncode != 0: + raise RuntimeError(f"ECR docker login failed: {(proc.stderr or '')[:1000]}") + + def _runtime_name(self) -> str: + # AgentCore runtime names allow [A-Za-z][A-Za-z0-9_]*. + safe = re.sub(r"[^a-zA-Z0-9_]+", "_", f"bf_{self.environment_name}") + return safe[:48] + + async def _ensure_runtime(self, image_uri: str) -> str: + """Reuse the agent runtime bound to *image_uri*, or create it.""" + control = self._client("bedrock-agentcore-control") + name = self._runtime_name() + + existing = await asyncio.to_thread(self._find_runtime, control, name) + if existing is not None: + arn, runtime_id, current_image = existing + self._runtime_id = runtime_id + if current_image == image_uri: + await asyncio.to_thread(self._wait_ready, control, runtime_id) + self.logger.info("Reusing AgentCore runtime %s", arn) + return arn + await asyncio.to_thread( + control.update_agent_runtime, + agentRuntimeId=runtime_id, + agentRuntimeArtifact={ + "containerConfiguration": {"containerUri": image_uri} + }, + roleArn=self._require_role_arn(), + networkConfiguration={"networkMode": "PUBLIC"}, + protocolConfiguration=_PROTOCOL_CONFIGURATION, + ) + await asyncio.to_thread(self._wait_ready, control, runtime_id) + self.logger.info("Updated AgentCore runtime %s to %s", arn, image_uri) + return arn + + created = await asyncio.to_thread( + control.create_agent_runtime, + agentRuntimeName=name, + agentRuntimeArtifact={ + "containerConfiguration": {"containerUri": image_uri} + }, + roleArn=self._require_role_arn(), + networkConfiguration={"networkMode": "PUBLIC"}, + protocolConfiguration=_PROTOCOL_CONFIGURATION, + lifecycleConfiguration=self._lifecycle_configuration(), + description=f"BenchFlow sandbox for {self.environment_name}", + ) + self._created_runtime = True + self._runtime_id = created["agentRuntimeId"] + await asyncio.to_thread(self._wait_ready, control, created["agentRuntimeId"]) + return created["agentRuntimeArn"] + + def _lifecycle_configuration(self) -> dict[str, int]: + """Session idle/lifetime caps. + + The service defaults (900 s idle, 8 h lifetime) are short enough that a + long agent turn can have its microVM reclaimed mid-run, which surfaces + as a dead transport rather than as an infrastructure error. Default the + idle window to the task's own agent timeout when that is larger. + """ + idle = int(os.environ.get(_ENV_IDLE_TIMEOUT) or 0) or max( + 900, int(getattr(self.task_env_config, "agent_timeout_sec", 0) or 0) + ) + lifetime = int(os.environ.get(_ENV_MAX_LIFETIME) or 0) or 28800 + return { + "idleRuntimeSessionTimeout": idle, + "maxLifetime": lifetime, + } + + def _require_role_arn(self) -> str: + if self._role_arn: + return self._role_arn + raise RuntimeError( + "AgentCore needs an execution role for the runtime. Set " + f"{_ENV_ROLE_ARN} to a role that bedrock-agentcore.amazonaws.com " + "can assume and that can pull from ECR and write CloudWatch logs." + ) + + @staticmethod + def _find_runtime(control: Any, name: str) -> tuple[str, str, str | None] | None: + paginator_pages = control.list_agent_runtimes() + for runtime in paginator_pages.get("agentRuntimes", []): + if runtime.get("agentRuntimeName") != name: + continue + runtime_id = runtime["agentRuntimeId"] + detail = control.get_agent_runtime(agentRuntimeId=runtime_id) + artifact = detail.get("agentRuntimeArtifact") or {} + image = (artifact.get("containerConfiguration") or {}).get("containerUri") + return detail["agentRuntimeArn"], runtime_id, image + return None + + @staticmethod + def _wait_ready(control: Any, runtime_id: str) -> None: + import time + + deadline = time.monotonic() + _RUNTIME_READY_TIMEOUT_SEC + status = "CREATING" + while time.monotonic() < deadline: + status = control.get_agent_runtime(agentRuntimeId=runtime_id)["status"] + if status == "READY": + return + if status in {"CREATE_FAILED", "UPDATE_FAILED"}: + raise SandboxStartupError( + f"AgentCore runtime entered {status}", + sandbox_id=runtime_id, + sandbox_state=status, + ) + time.sleep(_RUNTIME_POLL_INTERVAL_SEC) + raise SandboxStartupError( + f"AgentCore runtime not READY within {_RUNTIME_READY_TIMEOUT_SEC}s " + f"(last status {status})", + sandbox_id=runtime_id, + sandbox_state=status, + ) + + async def _warm_session(self) -> None: + """Boot the microVM and block until the command plane answers. + + ``READY`` describes the *runtime definition*, not a running session. + The first command on a new session is what actually pulls the image + and starts the container, and until the container is serving on 8080 + the service answers with a 500 ``RuntimeClientError``. Observed + directly: a cold, never-pulled image 500s on the first command and + succeeds moments later, while an image already warm in the region + succeeds immediately. + + Retrying here is what keeps that cold-start race from surfacing as a + task failure — the alternative is a rollout that scores 0 because the + sandbox was still booting. + """ + delay = _SESSION_WARMUP_BACKOFF_SEC + last: Exception | None = None + for attempt in range(1, _SESSION_WARMUP_ATTEMPTS + 1): + try: + result = await self.exec("true", timeout_sec=60) + if result.return_code == 0: + if attempt > 1: + self.logger.info( + "AgentCore session warm after %d attempts", attempt + ) + return + last = RuntimeError( + (result.stderr or result.stdout or "no output").strip()[:500] + ) + except SandboxStartupError: + # Throttling/quota — genuinely infrastructure, and retrying + # here would just burn the caller's time. + raise + except Exception as exc: + last = exc + if attempt < _SESSION_WARMUP_ATTEMPTS: + self.logger.debug( + "AgentCore session not warm yet (attempt %d/%d): %s", + attempt, + _SESSION_WARMUP_ATTEMPTS, + last, + ) + await asyncio.sleep(delay) + delay = min(delay * 2, _SESSION_WARMUP_MAX_BACKOFF_SEC) + + raise SandboxStartupError( + f"AgentCore session did not accept commands after " + f"{_SESSION_WARMUP_ATTEMPTS} attempts: {last}", + sandbox_id=self.runtime_arn, + ) + + async def stop(self, delete: bool) -> None: + await self._safe_teardown(delete_runtime=delete) + + async def _safe_teardown(self, *, delete_runtime: bool = True) -> None: + if self.runtime_arn and self.runtime_session_id: + try: + await asyncio.to_thread( + self._client("bedrock-agentcore").stop_runtime_session, + runtimeSessionId=self.runtime_session_id, + agentRuntimeArn=self.runtime_arn, + qualifier="DEFAULT", + ) + except Exception as exc: + self.logger.warning("Failed to stop AgentCore session: %s", exc) + self.runtime_session_id = None + + if delete_runtime and self._created_runtime and self._runtime_id: + try: + await asyncio.to_thread( + self._client("bedrock-agentcore-control").delete_agent_runtime, + agentRuntimeId=self._runtime_id, + ) + self.logger.info("Deleted AgentCore runtime %s", self._runtime_id) + except Exception as exc: + self.logger.warning("Failed to delete AgentCore runtime: %s", exc) + self._runtime_id = None + self.runtime_arn = None + self._created_runtime = False + + # ------------------------------------------------------------------ exec + + async def live_process(self, *, agent: str | None = None) -> LiveProcess: + from benchflow.sandbox.process import AgentCoreProcess + + return AgentCoreProcess.from_sandbox_env(self) + + def _reject_non_main(self, service: str) -> None: + if service != "main": + raise ValueError( + f"AgentCore is a single-container backend and cannot target " + f"service {service!r}. Multi-container (vulhub-style) tasks " + "require the Docker sandbox (#248)." + ) + + async def exec( + self, + command: str, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_sec: int | None = None, + user: str | int | None = None, + service: str = "main", + ) -> ExecResult: + self._reject_non_main(service) + if not self.runtime_arn or not self.runtime_session_id: + raise RuntimeError("AgentCore sandbox not started. Call start() first.") + + wrapped = command + merged_env = self._merge_env(env) + if merged_env: + wrapped = wrap_command_with_env_file( + merged_env, wrapped, env_path_prefix="/tmp/.bf_agentcore_env_" + ) + if cwd: + wrapped = f"cd {shlex.quote(cwd)} && {wrapped}" + + resolved_user = self._resolve_user(user) + if resolved_user is not None: + if isinstance(resolved_user, int): + user_arg = f"$(getent passwd {resolved_user} | cut -d: -f1)" + else: + user_arg = shlex.quote(str(resolved_user)) + wrapped = f"su {user_arg} -s /bin/bash -c {shlex.quote(wrapped)}" + + payload = f"/bin/bash -c {shlex.quote(wrapped)}" + # The service enforces 1..3600 for the command timeout. + timeout = max(1, min(int(timeout_sec or 300), 3600)) + return await asyncio.to_thread(self._invoke_command, payload, timeout) + + def _invoke_command(self, payload: str, timeout: int) -> ExecResult: + client = self._client("bedrock-agentcore") + response = client.invoke_agent_runtime_command( + agentRuntimeArn=self.runtime_arn, + runtimeSessionId=self.runtime_session_id, + qualifier="DEFAULT", + contentType="application/json", + accept="application/vnd.amazon.eventstream", + body={"command": payload, "timeout": timeout}, + ) + + stdout: list[str] = [] + stderr: list[str] = [] + exit_code: int | None = None + status: str | None = None + + for event in response.get("stream", []): + chunk = event.get("chunk") + if chunk is None: + self._raise_for_stream_error(event) + continue + delta = chunk.get("contentDelta") + if delta: + stdout.append(delta.get("stdout") or "") + stderr.append(delta.get("stderr") or "") + stop = chunk.get("contentStop") + if stop: + exit_code = stop.get("exitCode") + status = stop.get("status") + + if status == "TIMED_OUT": + raise TimeoutError( + f"AgentCore command timed out after {timeout}s: {payload[:200]}" + ) + return ExecResult( + stdout="".join(stdout), + stderr="".join(stderr), + # A stream that stops without contentStop has no exit code to + # report; surfacing 1 keeps callers' `return_code != 0` checks + # honest rather than inventing a success. + return_code=exit_code if exit_code is not None else 1, + ) + + @staticmethod + def _raise_for_stream_error(event: dict[str, Any]) -> None: + """Translate a typed stream error into the right BenchFlow failure. + + Throttling and quota exhaustion are infrastructure, not agent + behaviour; raising ``SandboxStartupError`` keeps them attributable in + ``result.json`` instead of being recorded as a failed command. + """ + infra = { + "throttlingException", + "serviceQuotaExceededException", + "internalServerException", + } + for key, value in event.items(): + message = ( + (value or {}).get("message", "") if isinstance(value, dict) else "" + ) + if key in infra: + raise SandboxStartupError( + f"AgentCore {key}: {message}", sandbox_state=key + ) + if key.endswith("Exception") or key == "runtimeClientError": + raise RuntimeError(f"AgentCore {key}: {message}") + + # -------------------------------------------------------- file transfer + + async def write_text_file( + self, remote_path: str, body: str, *, mode: str = "600" + ) -> bool: + """Write *body* to *remote_path* inside the session, base64-encoded. + + Used by the ACP transport to stage the agent env file without typing + secrets into the PTY, where they would be echoed into the agent log. + """ + encoded = base64.b64encode(body.encode()).decode() + if len(encoded) > _MAX_INLINE_UPLOAD_BYTES: + raise ValueError( + f"Refusing to inline {len(encoded)} bytes into an AgentCore " + f"command (limit {_MAX_INLINE_UPLOAD_BYTES}). " + ) + quoted = shlex.quote(remote_path) + result = await self.exec( + f"mkdir -p $(dirname {quoted}) && " + f"printf %s {shlex.quote(encoded)} | base64 -d > {quoted} && " + f"chmod {mode} {quoted}", + timeout_sec=60, + ) + return result.return_code == 0 + + async def upload_file(self, source_path: Path | str, target_path: str) -> None: + source = Path(source_path) + data = source.read_bytes() + encoded = base64.b64encode(data).decode() + if len(encoded) > _MAX_INLINE_UPLOAD_BYTES: + # Larger payloads go through the tar path, which streams in + # bounded chunks instead of one oversized command. + await self._upload_via_tar( + {source: PurePosixPath(target_path)}, + root=source.parent, + ) + return + quoted = shlex.quote(target_path) + result = await self.exec( + f"mkdir -p $(dirname {quoted}) && " + f"printf %s {shlex.quote(encoded)} | base64 -d > {quoted}", + timeout_sec=120, + ) + if result.return_code != 0: + raise RuntimeError( + f"AgentCore upload_file failed: {(result.stderr or '')[:500]}" + ) + + async def upload_dir( + self, source_dir: Path | str, target_dir: str, service: str = "main" + ) -> None: + self._reject_non_main(service) + source = Path(source_dir) + if not source.is_dir(): + raise FileNotFoundError(f"Source directory {source} does not exist") + # Skip symlinks (#411): a task-controlled link must not exfiltrate host + # files into the remote microVM. + members = { + path: PurePosixPath(target_dir) / path.relative_to(source).as_posix() + for path in iter_safe_tree(source, context=f"agentcore upload_dir {source}") + } + await self.exec(f"mkdir -p {shlex.quote(target_dir)}", user="root") + if members: + await self._upload_via_tar(members, root=source) + + async def _upload_via_tar( + self, members: dict[Path, PurePosixPath], root: Path + ) -> None: + """Ship files as a single base64 tar stream, chunked under the cap. + + One archive per directory keeps this O(1) commands for the common case + instead of one round trip per file, and the chunk loop keeps every + individual command well inside the service's 64 KB payload limit. + """ + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w:gz") as tar: + for local, remote in members.items(): + tar.add(local, arcname=str(remote).lstrip("/")) + encoded = base64.b64encode(buffer.getvalue()).decode() + + staging = f"/tmp/.bf_upload_{uuid.uuid4().hex[:16]}.b64" + chunk = _MAX_INLINE_UPLOAD_BYTES + for index in range(0, len(encoded), chunk): + piece = encoded[index : index + chunk] + redirect = ">" if index == 0 else ">>" + result = await self.exec( + f"printf %s {shlex.quote(piece)} {redirect} {staging}", + timeout_sec=120, + ) + if result.return_code != 0: + await self.exec(f"rm -f {staging}", timeout_sec=30) + raise RuntimeError( + f"AgentCore upload staging failed: {(result.stderr or '')[:500]}" + ) + + result = await self.exec( + f"base64 -d {staging} | tar -xzf - -C / && rm -f {staging}", + timeout_sec=600, + user="root", + ) + if result.return_code != 0: + raise RuntimeError( + f"AgentCore upload extract failed: {(result.stderr or '')[:500]}" + ) + + async def download_file(self, source_path: str, target_path: Path | str) -> None: + target = Path(target_path) + target.parent.mkdir(parents=True, exist_ok=True) + result = await self.exec(f"base64 {shlex.quote(source_path)}", timeout_sec=300) + if result.return_code != 0: + raise RuntimeError( + f"AgentCore download_file failed: {(result.stderr or '')[:500]}" + ) + target.write_bytes(base64.b64decode((result.stdout or "").replace("\n", ""))) + + async def download_dir( + self, source_dir: str, target_dir: Path | str, service: str = "main" + ) -> None: + self._reject_non_main(service) + target = Path(target_dir) + target.mkdir(parents=True, exist_ok=True) + result = await self.exec( + f"tar -czf - -C {shlex.quote(source_dir)} . | base64 -w0", + timeout_sec=600, + ) + if result.return_code != 0: + raise RuntimeError( + f"AgentCore download_dir failed: {(result.stderr or '')[:500]}" + ) + raw = (result.stdout or "").strip() + if not raw: + return + blob = base64.b64decode(raw) + if len(blob) > _MAX_DOWNLOAD_BYTES: + raise RuntimeError( + f"AgentCore download_dir payload {len(blob)} bytes exceeds the " + f"{_MAX_DOWNLOAD_BYTES} byte cap; narrow the directory." + ) + with tarfile.open(fileobj=io.BytesIO(blob), mode="r:gz") as tar: + tar.extractall(target, filter="data") diff --git a/src/benchflow/sandbox/apple_container.py b/src/benchflow/sandbox/apple_container.py index 0693c5de2..65e94ab25 100644 --- a/src/benchflow/sandbox/apple_container.py +++ b/src/benchflow/sandbox/apple_container.py @@ -17,9 +17,13 @@ import subprocess import sys from pathlib import Path, PurePosixPath +from typing import TYPE_CHECKING from benchflow.sandbox._base import BaseSandbox, ExecResult, wrap_command_with_env_file +if TYPE_CHECKING: + from benchflow.sandbox.process import LiveProcess + _MIN_CONTAINER_VERSION = (1, 1, 0) _KALLOC_SAFE_LIMIT = 3_000_000 _KALLOC_MIN_HEADROOM = 200_000 @@ -168,6 +172,11 @@ def is_mounted(self) -> bool: def sandbox_id(self) -> str | None: return self._container_name + async def live_process(self, *, agent: str | None = None) -> LiveProcess: + from benchflow.sandbox.process import AppleContainerProcess + + return AppleContainerProcess.from_sandbox_env(self) + @classmethod def preflight(cls) -> None: if sys.platform != "darwin": diff --git a/src/benchflow/sandbox/daytona.py b/src/benchflow/sandbox/daytona.py index da7cc4bb7..6ca7638da 100644 --- a/src/benchflow/sandbox/daytona.py +++ b/src/benchflow/sandbox/daytona.py @@ -14,7 +14,7 @@ import logging import shlex from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 try: @@ -85,6 +85,9 @@ def wait_exponential(*_args: Any, **_kwargs: Any) -> Any: from benchflow.task.config import SandboxConfig from benchflow.task.paths import RolloutPaths +if TYPE_CHECKING: + from benchflow.sandbox.process import LiveProcess + # ``SandboxStartupError`` used to live in this module. It now lives in # ``benchflow.sandbox.protocol`` so a base install without the # ``sandbox-daytona`` extra can still import ``benchflow.rollout`` (issue #358). @@ -909,6 +912,23 @@ async def is_file( return await super().is_file(path, user=user, service=service) return await self._strategy.is_file(path, user=self._resolve_user(user)) + async def live_process(self, *, agent: str | None = None) -> LiveProcess: + """Pick the Daytona ACP transport (PTY by default, SSH for gemini). + + Selection lives here rather than at the ACP layer so the choice and the + provenance string reported by ``selected_acp_transport`` stay derived + from the same helper. + """ + from benchflow.acp.selection import selected_acp_transport + from benchflow.sandbox.process import DaytonaProcess, DaytonaPtyProcess + + transport = selected_acp_transport(agent=agent or "", environment="daytona") + if transport == "ssh": + logger.info("Using SSH transport for %s on Daytona", agent) + return await DaytonaProcess.from_sandbox_env(self) + logger.info("Using PTY transport for Daytona sandbox") + return await DaytonaPtyProcess.from_sandbox_env(self) + async def attach(self) -> None: return await self._strategy.attach() diff --git a/src/benchflow/sandbox/docker.py b/src/benchflow/sandbox/docker.py index 1c81a581b..c14b18fa9 100644 --- a/src/benchflow/sandbox/docker.py +++ b/src/benchflow/sandbox/docker.py @@ -18,7 +18,7 @@ import sys import uuid from pathlib import Path -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar from pydantic import BaseModel @@ -41,6 +41,9 @@ from benchflow.task.env import resolve_env_vars from benchflow.task.paths import RolloutPaths, SandboxPaths +if TYPE_CHECKING: + from benchflow.sandbox.process import LiveProcess + logger = logging.getLogger("benchflow") _DOCKER_BUILD_RETRY_DELAYS_SEC = (2.0, 5.0) @@ -867,6 +870,11 @@ def _wrap_command_with_env_file(cls, env: dict[str, str], command: str) -> str: env, command, env_path_prefix=cls._ENV_FILE_PREFIX ) + async def live_process(self, *, agent: str | None = None) -> LiveProcess: + from benchflow.sandbox.process import DockerProcess + + return DockerProcess.from_sandbox_env(self) + async def attach(self) -> None: variables = " ".join( f"export {k}={shlex.quote(str(v))}" diff --git a/src/benchflow/sandbox/process/__init__.py b/src/benchflow/sandbox/process/__init__.py new file mode 100644 index 000000000..6677813dc --- /dev/null +++ b/src/benchflow/sandbox/process/__init__.py @@ -0,0 +1,39 @@ +"""Live stdio connections to a process inside a sandbox. + +Provides a bidirectional pipe (send lines in, read lines out) needed for +ACP agents running inside containers. Implementations: + +- DockerProcess: `docker compose exec -i` (local Docker) +- AppleContainerProcess: `container exec -i` (Apple Container) +- DaytonaProcess: SSH to a Daytona sandbox +- DaytonaPtyProcess: Daytona PTY WebSocket +- AgentCoreProcess: Bedrock AgentCore shell WebSocket + +The first three are local-subprocess transports and share +:class:`SubprocessLiveProcess`; the last two carry bytes over a WebSocket and +implement :class:`LiveProcess` directly. + +This package was a single ``process.py`` until it outgrew 1000 lines. Import +paths are unchanged — everything public is re-exported here. +""" + +from benchflow.sandbox.process._base import ( + LiveProcess, + SubprocessLiveProcess, + drain_oversized_line, +) +from benchflow.sandbox.process.agentcore import AgentCoreProcess +from benchflow.sandbox.process.apple import AppleContainerProcess +from benchflow.sandbox.process.daytona import DaytonaProcess, DaytonaPtyProcess +from benchflow.sandbox.process.docker import DockerProcess + +__all__ = [ + "AgentCoreProcess", + "AppleContainerProcess", + "DaytonaProcess", + "DaytonaPtyProcess", + "DockerProcess", + "LiveProcess", + "SubprocessLiveProcess", + "drain_oversized_line", +] diff --git a/src/benchflow/sandbox/process/_base.py b/src/benchflow/sandbox/process/_base.py new file mode 100644 index 000000000..35b409351 --- /dev/null +++ b/src/benchflow/sandbox/process/_base.py @@ -0,0 +1,193 @@ +"""Transport-agnostic ``LiveProcess`` contract and its subprocess implementation. + +``LiveProcess`` is the bidirectional line pipe an ACP agent speaks over. It +deliberately declares *only* the four operations the ACP transport needs +(``start``/``readline``/``writeline``/``close``) plus ``is_running``, with no +assumption about what carries the bytes. + +:class:`SubprocessLiveProcess` supplies the local-``asyncio``-subprocess +implementation shared by the Docker, Apple Container, and Daytona-SSH +backends: for those three the pipe *is* a child process's stdio. + +The split exists because two transports are not subprocess-backed at all — +``DaytonaPtyProcess`` (Daytona PTY WebSocket) and ``AgentCoreProcess`` +(Bedrock AgentCore shell WebSocket). Before the split they inherited +subprocess semantics they could not honour and had to neutralize the base +class with ``_process = None # Not used`` plus a full override of +``readline``/``writeline``/``close``. One such escape hatch is a wart; two +means the base class was modelling the wrong thing. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import re +from abc import ABC, abstractmethod + +logger = logging.getLogger(__name__) + +_BUFFER_LIMIT = 10 * 1024 * 1024 # 10MB readline buffer +_DIAG_TRUNCATE = 2000 # max chars for diagnostic stderr in error messages +_BOOTSTRAP_DONE = "__BENCHFLOW_BOOTSTRAP_DONE__" +_ENV_KEY_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +async def drain_oversized_line(reader: asyncio.StreamReader) -> int: + """Drain an oversized line from *reader* after a buffer overflow. + + Clears the internal buffer and attempts to skip ahead to the next + newline. Returns the number of bytes discarded. + """ + # Reach into asyncio.StreamReader internals to clear the buffer after + # a LimitOverrunError. There's no public API for this; the private + # attributes are stable across Python 3.10+. + skipped = len(reader._buffer) # ty: ignore[unresolved-attribute] + reader._buffer.clear() # ty: ignore[unresolved-attribute] + reader._maybe_resume_transport() # ty: ignore[unresolved-attribute] + try: + await asyncio.wait_for(reader.readuntil(b"\n"), timeout=5) + except Exception: + logger.debug("Could not find next newline after buffer overflow") + return skipped + + +class LiveProcess(ABC): + """Abstract live stdin/stdout connection to a process inside a sandbox. + + Implementations carry the bytes however their backend allows — a local + child process's stdio, an SSH pipe, or a WebSocket terminal. Nothing in + this contract presumes a subprocess; backends that *are* subprocess-backed + should extend :class:`SubprocessLiveProcess` instead of implementing the + read/write/close trio by hand. + """ + + @abstractmethod + async def start( + self, + command: str, + env: dict[str, str] | None = None, + cwd: str | None = None, + ) -> None: + """Start the process with live stdin/stdout.""" + + @abstractmethod + async def readline(self) -> bytes: + """Read one line from the process's stdout.""" + + @abstractmethod + async def writeline(self, data: str) -> None: + """Write one line to the process's stdin.""" + + @abstractmethod + async def close(self) -> None: + """Terminate the process (idempotent — safe to call after death).""" + + @property + @abstractmethod + def is_running(self) -> bool: + """Whether the transport is still usable.""" + + +class SubprocessLiveProcess(LiveProcess): + """A :class:`LiveProcess` whose pipe is a local ``asyncio`` subprocess. + + Subclasses only implement ``start`` — spawning whatever CLI reaches into + their sandbox (``docker compose exec -i``, ``container exec -i``, ``ssh``) + and assigning the result to ``self._process``. Everything else is shared. + """ + + _process: asyncio.subprocess.Process | None = None + + async def readline(self) -> bytes: + """Read one line from stdout.""" + if not self._process or not self._process.stdout: + raise RuntimeError("Process not started") + try: + line = await self._process.stdout.readline() + except (ValueError, asyncio.LimitOverrunError) as e: + # Buffer overflow — line exceeds _BUFFER_LIMIT. + skipped = await drain_oversized_line(self._process.stdout) + logger.warning(f"Skipped oversized line ({skipped} bytes): {e}") + # Return empty line — caller will retry readline + return b"" + if not line: + stderr_text = "" + if self._process and self._process.stderr: + try: + stderr_bytes = await asyncio.wait_for( + self._process.stderr.read(8192), timeout=2 + ) + stderr_text = stderr_bytes.decode(errors="replace").strip() + except Exception: + logger.debug("Could not read stderr from closed process") + rc = self._process.returncode if self._process else None + # Diagnose: rc=None with closed stdout usually means the *transport* + # died (SSH/Daytona idle sleep, container killed) while the local + # subprocess wrapper is still alive. rc set means the local process + # actually exited. Surfacing the distinction makes the failure + # actionable instead of cryptic. + pid = self._process.pid if self._process else None + if rc is None: + hint = ( + f"Local subprocess (pid={pid}) is still alive but its " + "stdout/transport closed. This usually means the remote " + "container or SSH session was killed (e.g. Daytona idle " + "sleep, agent hung with no output)." + ) + diagnosis = "remote_session_killed" + else: + hint = f"Local subprocess exited with rc={rc} before stdout closed." + diagnosis = "process_exited" + msg = f"Process closed stdout (rc={rc}): {hint}" + stderr_snippet: str | None = None + if stderr_text: + stderr_snippet = stderr_text[:_DIAG_TRUNCATE] + msg += f"\nstderr: {stderr_snippet}" + # Raise a structured TransportClosedError at the source so + # downstream code (rollout._build_rollout_result) doesn't have + # to regex-parse the human-readable message back into fields + # (issue #504). + from benchflow.diagnostics import ( + TransportClosedDiagnostic, + TransportClosedError, + ) + + raise TransportClosedError( + msg, + TransportClosedDiagnostic( + raw_message=msg[:500], + process_exit_code=rc, + process_pid=pid, + transport_diagnosis=diagnosis, + stderr_snippet=stderr_snippet, + ), + ) + return line + + async def writeline(self, data: str) -> None: + """Write one line to stdin.""" + if not self._process or not self._process.stdin: + raise RuntimeError("Process not started") + self._process.stdin.write((data + "\n").encode()) + await self._process.stdin.drain() + + async def close(self) -> None: + """Terminate the process (idempotent — safe to call after process death).""" + if self._process: + if self._process.stdin: + with contextlib.suppress(OSError): # already closed + self._process.stdin.close() + if self._process.returncode is None: + self._process.terminate() + try: + await asyncio.wait_for(self._process.wait(), timeout=5) + except TimeoutError: + self._process.kill() + await self._process.wait() + logger.info("Process terminated") + + @property + def is_running(self) -> bool: + return self._process is not None and self._process.returncode is None diff --git a/src/benchflow/sandbox/process/agentcore.py b/src/benchflow/sandbox/process/agentcore.py new file mode 100644 index 000000000..2e78cdcfb --- /dev/null +++ b/src/benchflow/sandbox/process/agentcore.py @@ -0,0 +1,260 @@ +"""Live stdio through the Bedrock AgentCore interactive shell WebSocket. + +``InvokeAgentRuntimeCommand`` (used by ``AgentCoreSandbox.exec``) is one-shot: +each call spawns a fresh bash, runs it to completion, and returns. It cannot +hold the long-lived bidirectional pipe an ACP agent speaks JSON-RPC over. That +is what ``open_shell`` provides — a persistent WebSocket terminal attached to +the *same* runtime session, so the agent started here shares a filesystem with +every ``exec()`` the kernel and verifier make. + +The channel is a **PTY**, not a raw pipe: it echoes input, emits bracketed-paste +control sequences, and terminates lines with CRLF. That is the same shape as +``DaytonaPtyProcess``, and this class handles it the same proven way — put the +line discipline into raw/no-echo mode, synchronize on a nonce marker before +handing the terminal to the agent, and strip CR while framing lines. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import shlex +import uuid +from typing import Any + +from benchflow.sandbox.process._base import LiveProcess + +logger = logging.getLogger(__name__) + +_START_MARKER_TIMEOUT_SEC = 180 +_READLINE_TIMEOUT_ENV = "BENCHFLOW_AGENTCORE_READLINE_TIMEOUT" +_READLINE_TIMEOUT_DEFAULT_SEC = 900.0 + + +def _readline_timeout_sec() -> float: + import os + + value = os.environ.get(_READLINE_TIMEOUT_ENV) + if value is None: + return _READLINE_TIMEOUT_DEFAULT_SEC + try: + timeout = float(value) + except ValueError: + logger.warning( + "Invalid %s=%r; using default %.0fs", + _READLINE_TIMEOUT_ENV, + value, + _READLINE_TIMEOUT_DEFAULT_SEC, + ) + return _READLINE_TIMEOUT_DEFAULT_SEC + return timeout if timeout > 0 else _READLINE_TIMEOUT_DEFAULT_SEC + + +class AgentCoreProcess(LiveProcess): + """Live stdin/stdout over an AgentCore ``open_shell`` WebSocket terminal.""" + + def __init__( + self, + sandbox: Any, + runtime_arn: str, + session_id: str, + region: str, + ) -> None: + self._sandbox = sandbox + self._runtime_arn = runtime_arn + self._session_id = session_id + self._region = region + self._shell: Any = None + self._reader_task: asyncio.Task[None] | None = None + self._line_buffer: asyncio.Queue[bytes] = asyncio.Queue() + self._partial = b"" + self._closed = False + self._failure: BaseException | None = None + + @classmethod + def from_sandbox_env(cls, env: Any) -> AgentCoreProcess: + """Create from a started :class:`AgentCoreSandbox`.""" + runtime_arn = getattr(env, "runtime_arn", None) + session_id = getattr(env, "runtime_session_id", None) + if not runtime_arn or not session_id: + raise RuntimeError("AgentCore sandbox not started") + return cls( + sandbox=env, + runtime_arn=runtime_arn, + session_id=session_id, + region=env.region, + ) + + async def _drain_frames(self) -> None: + """Frame WebSocket payloads into newline-terminated lines.""" + try: + async for frame in self._shell: + payload = frame.payload + if isinstance(payload, str): + payload = payload.encode() + if not payload: + continue + self._partial += payload + while b"\n" in self._partial: + line, self._partial = self._partial.split(b"\n", 1) + await self._line_buffer.put(line.replace(b"\r", b"") + b"\n") + except asyncio.CancelledError: + raise + except BaseException as exc: + # Record rather than raise: this runs in a background task whose + # exception would otherwise be swallowed, leaving readline() to + # hang until its timeout instead of reporting the real cause. + self._failure = exc + logger.warning("AgentCore shell reader stopped: %s", exc) + + async def _write_env_file(self, env: dict[str, str]) -> str: + """Materialize *env* as a mode-0600 file inside the session container. + + Routed through the sandbox's own ``exec`` (and therefore through the + canonical base64 env-file wrapper) so secrets never appear as literal + text typed into the terminal, where the PTY would echo them straight + back into the agent log. + """ + remote_path = f"/tmp/.benchflow_agent_env_{uuid.uuid4().hex[:16]}" + body = "".join(f"export {k}={shlex.quote(v)}\n" for k, v in env.items()) + result = await self._sandbox.write_text_file(remote_path, body, mode="600") + if result is False: + raise RuntimeError("Failed to stage AgentCore agent env file") + return remote_path + + async def start( + self, + command: str, + env: dict[str, str] | None = None, + cwd: str | None = None, + ) -> None: + from bedrock_agentcore.runtime import AgentCoreRuntimeClient + + parts: list[str] = [] + if cwd: + parts.append(f"cd {shlex.quote(cwd)}") + if env: + remote_env_path = await self._write_env_file(env) + quoted = shlex.quote(remote_env_path) + parts.append(f". {quoted}") + parts.append(f"rm -f {quoted}") + parts.append(f"exec bash -lc {shlex.quote(command)}") + launch = " && ".join(parts) + + client = AgentCoreRuntimeClient(region=self._region) + shell = client.open_shell( + runtime_arn=self._runtime_arn, + session_id=self._session_id, + ) + try: + self._shell = await shell.__aenter__() + logger.info( + "AgentCore shell connected (shell_id=%s, session=%s)", + getattr(self._shell, "shell_id", "?"), + self._session_id, + ) + self._reader_task = asyncio.create_task(self._drain_frames()) + + # Take the terminal out of cooked mode before any ACP traffic. + # ACP JSON-RPC frames routinely exceed the 4096-byte canonical-mode + # line limit, and echo would feed the agent's own output back to it. + marker = f"__BENCHFLOW_ACP_{uuid.uuid4().hex[:12]}__" + await self._shell.send( + "stty raw -echo 2>/dev/null || " + "stty -echo -icanon min 1 time 0 2>/dev/null || true; " + f"echo '{marker}'\n" + ) + await self._await_marker(marker) + self._clear_buffered_output() + await self._shell.send(launch + "\n") + except BaseException: + await self.close() + raise + logger.info("AgentCore shell marker seen, agent starting") + + async def _await_marker(self, marker: str) -> None: + from benchflow.diagnostics import ( + TransportClosedDiagnostic, + TransportClosedError, + ) + + while True: + try: + line = await asyncio.wait_for( + self._line_buffer.get(), timeout=_START_MARKER_TIMEOUT_SEC + ) + except TimeoutError as e: + msg = ( + "AgentCore shell: timed out waiting for the start marker " + f"(session={self._session_id})" + ) + raise TransportClosedError( + msg, + TransportClosedDiagnostic( + raw_message=msg, + transport_diagnosis="pty_startup_timeout", + ), + ) from e + if marker in line.decode(errors="replace"): + return + + def _clear_buffered_output(self) -> None: + self._partial = b"" + while not self._line_buffer.empty(): + with contextlib.suppress(asyncio.QueueEmpty): + self._line_buffer.get_nowait() + + async def readline(self) -> bytes: + from benchflow.diagnostics import ( + TransportClosedDiagnostic, + TransportClosedError, + ) + + def _closed(msg: str, diagnosis: str) -> TransportClosedError: + return TransportClosedError( + msg, + TransportClosedDiagnostic( + raw_message=msg[:500], transport_diagnosis=diagnosis + ), + ) + + if self._closed: + raise _closed("AgentCore shell closed", "pty_error") + timeout = _readline_timeout_sec() + try: + return await asyncio.wait_for(self._line_buffer.get(), timeout=timeout) + except TimeoutError as e: + # A dead reader task means the WebSocket dropped; report that rather + # than the timeout, which would misattribute an infra failure to a + # silent agent. + if self._failure is not None: + raise _closed( + f"AgentCore shell transport failed: {self._failure}", + "remote_session_killed", + ) from self._failure + raise _closed( + f"AgentCore shell readline timeout ({timeout:g}s)", "pty_error" + ) from e + + async def writeline(self, data: str) -> None: + if not self._shell or self._closed: + raise RuntimeError("AgentCore shell not started") + await self._shell.send(data + "\n") + + async def close(self) -> None: + self._closed = True + if self._reader_task is not None: + self._reader_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._reader_task + self._reader_task = None + if self._shell is not None: + with contextlib.suppress(Exception): + await self._shell.close() + self._shell = None + logger.info("AgentCore shell terminated") + + @property + def is_running(self) -> bool: + return self._shell is not None and not self._closed diff --git a/src/benchflow/sandbox/process/apple.py b/src/benchflow/sandbox/process/apple.py new file mode 100644 index 000000000..8d1311ba9 --- /dev/null +++ b/src/benchflow/sandbox/process/apple.py @@ -0,0 +1,101 @@ +"""Live stdio through Apple Container's native exec transport.""" + +from __future__ import annotations + +import asyncio +import logging +import shlex +import uuid +from typing import Any + +from benchflow.sandbox.process._base import ( + _BUFFER_LIMIT, + _ENV_KEY_RE, + SubprocessLiveProcess, +) + +logger = logging.getLogger(__name__) + + +class AppleContainerProcess(SubprocessLiveProcess): + """Live stdin/stdout through Apple Container's native exec transport.""" + + def __init__(self, container_name: str): + self._container_name = container_name + self._env_path = f"/tmp/.benchflow_agent_env_{uuid.uuid4().hex[:16]}" + + @classmethod + def from_sandbox_env(cls, env: Any) -> AppleContainerProcess: + """Create from a started AppleContainerSandbox.""" + + container_name = getattr(env, "_container_name", None) + if not isinstance(container_name, str) or not container_name: + raise RuntimeError("Apple Container sandbox not started") + return cls(container_name) + + async def _write_env_to_container(self, env: dict[str, str]) -> None: + invalid = [key for key in env if not _ENV_KEY_RE.match(key)] + if invalid: + raise ValueError( + "Invalid environment variable name(s): " + ", ".join(sorted(invalid)) + ) + + lines = "".join( + f"export {key}={shlex.quote(value)}\n" for key, value in env.items() + ) + env_path = shlex.quote(self._env_path) + proc = await asyncio.create_subprocess_exec( + "container", + "exec", + "--interactive", + "--user", + "root", + self._container_name, + "sh", + "-c", + f"cat > {env_path} && chmod 600 {env_path}", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + _, stderr = await asyncio.wait_for( + proc.communicate(lines.encode()), timeout=30 + ) + except TimeoutError: + proc.kill() + await proc.wait() + raise + if proc.returncode != 0: + raise RuntimeError( + "Failed to write agent env in Apple container " + f"(rc={proc.returncode}): {stderr.decode(errors='replace')[:500]}" + ) + + async def start( + self, + command: str, + env: dict[str, str] | None = None, + cwd: str | None = None, + ) -> None: + if env: + await self._write_env_to_container(env) + env_path = shlex.quote(self._env_path) + command = f". {env_path} && rm -f {env_path} && {command}" + + args = ["container", "exec", "--interactive"] + if cwd: + args.extend(["--workdir", cwd]) + args.extend([self._container_name, "bash", "-c", command]) + self._process = await asyncio.create_subprocess_exec( + *args, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + limit=_BUFFER_LIMIT, + ) + logger.info( + "Apple Container process started (pid=%s, container=%s)", + self._process.pid, + self._container_name, + ) diff --git a/src/benchflow/sandbox/process.py b/src/benchflow/sandbox/process/daytona.py similarity index 63% rename from src/benchflow/sandbox/process.py rename to src/benchflow/sandbox/process/daytona.py index 3d7b6d2af..03931a426 100644 --- a/src/benchflow/sandbox/process.py +++ b/src/benchflow/sandbox/process/daytona.py @@ -1,30 +1,27 @@ -"""Live stdio connection to a process inside a sandbox. +"""Live stdio to a Daytona sandbox — SSH subprocess and PTY WebSocket variants.""" -Provides a bidirectional pipe (send lines in, read lines out) needed for -ACP agents running inside containers. Implementations: - -- DockerProcess: uses `docker compose exec -i` (local Docker) -- AppleContainerProcess: uses `container exec -i` (Apple Container) -- DaytonaProcess: uses SSH to a Daytona sandbox -""" +from __future__ import annotations import asyncio import contextlib import logging import os -import re import shlex import tempfile import uuid -from abc import ABC, abstractmethod from typing import Any +from benchflow.sandbox.process._base import ( + _BOOTSTRAP_DONE, + _BUFFER_LIMIT, + _DIAG_TRUNCATE, + _ENV_KEY_RE, + LiveProcess, + SubprocessLiveProcess, +) + logger = logging.getLogger(__name__) -_BUFFER_LIMIT = 10 * 1024 * 1024 # 10MB readline buffer -_DIAG_TRUNCATE = 2000 # max chars for diagnostic stderr in error messages -_BOOTSTRAP_DONE = "__BENCHFLOW_BOOTSTRAP_DONE__" -_ENV_KEY_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") _DAYTONA_PTY_READLINE_TIMEOUT_ENV = "BENCHFLOW_DAYTONA_PTY_READLINE_TIMEOUT" _DAYTONA_PTY_READLINE_TIMEOUT_DEFAULT_SEC = 900.0 _DAYTONA_SSH_ACCESS_TTL_MINUTES = 48 * 60 @@ -120,361 +117,7 @@ async def _bootstrap_daytona_env_file( ) -async def drain_oversized_line(reader: asyncio.StreamReader) -> int: - """Drain an oversized line from *reader* after a buffer overflow. - - Clears the internal buffer and attempts to skip ahead to the next - newline. Returns the number of bytes discarded. - """ - # Reach into asyncio.StreamReader internals to clear the buffer after - # a LimitOverrunError. There's no public API for this; the private - # attributes are stable across Python 3.10+. - skipped = len(reader._buffer) # ty: ignore[unresolved-attribute] - reader._buffer.clear() # ty: ignore[unresolved-attribute] - reader._maybe_resume_transport() # ty: ignore[unresolved-attribute] - try: - await asyncio.wait_for(reader.readuntil(b"\n"), timeout=5) - except Exception: - logger.debug("Could not find next newline after buffer overflow") - return skipped - - -class LiveProcess(ABC): - """Abstract live stdin/stdout connection to a process inside a sandbox.""" - - _process: asyncio.subprocess.Process | None = None - - @abstractmethod - async def start( - self, - command: str, - env: dict[str, str] | None = None, - cwd: str | None = None, - ) -> None: - """Start the process with live stdin/stdout.""" - - async def readline(self) -> bytes: - """Read one line from stdout.""" - if not self._process or not self._process.stdout: - raise RuntimeError("Process not started") - try: - line = await self._process.stdout.readline() - except (ValueError, asyncio.LimitOverrunError) as e: - # Buffer overflow — line exceeds _BUFFER_LIMIT. - skipped = await drain_oversized_line(self._process.stdout) - logger.warning(f"Skipped oversized line ({skipped} bytes): {e}") - # Return empty line — caller will retry readline - return b"" - if not line: - stderr_text = "" - if self._process and self._process.stderr: - try: - stderr_bytes = await asyncio.wait_for( - self._process.stderr.read(8192), timeout=2 - ) - stderr_text = stderr_bytes.decode(errors="replace").strip() - except Exception: - logger.debug("Could not read stderr from closed process") - rc = self._process.returncode if self._process else None - # Diagnose: rc=None with closed stdout usually means the *transport* - # died (SSH/Daytona idle sleep, container killed) while the local - # subprocess wrapper is still alive. rc set means the local process - # actually exited. Surfacing the distinction makes the failure - # actionable instead of cryptic. - pid = self._process.pid if self._process else None - if rc is None: - hint = ( - f"Local subprocess (pid={pid}) is still alive but its " - "stdout/transport closed. This usually means the remote " - "container or SSH session was killed (e.g. Daytona idle " - "sleep, agent hung with no output)." - ) - diagnosis = "remote_session_killed" - else: - hint = f"Local subprocess exited with rc={rc} before stdout closed." - diagnosis = "process_exited" - msg = f"Process closed stdout (rc={rc}): {hint}" - stderr_snippet: str | None = None - if stderr_text: - stderr_snippet = stderr_text[:_DIAG_TRUNCATE] - msg += f"\nstderr: {stderr_snippet}" - # Raise a structured TransportClosedError at the source so - # downstream code (rollout._build_rollout_result) doesn't have - # to regex-parse the human-readable message back into fields - # (issue #504). - from benchflow.diagnostics import ( - TransportClosedDiagnostic, - TransportClosedError, - ) - - raise TransportClosedError( - msg, - TransportClosedDiagnostic( - raw_message=msg[:500], - process_exit_code=rc, - process_pid=pid, - transport_diagnosis=diagnosis, - stderr_snippet=stderr_snippet, - ), - ) - return line - - async def writeline(self, data: str) -> None: - """Write one line to stdin.""" - if not self._process or not self._process.stdin: - raise RuntimeError("Process not started") - self._process.stdin.write((data + "\n").encode()) - await self._process.stdin.drain() - - async def close(self) -> None: - """Terminate the process (idempotent — safe to call after process death).""" - if self._process: - if self._process.stdin: - with contextlib.suppress(OSError): # already closed - self._process.stdin.close() - if self._process.returncode is None: - self._process.terminate() - try: - await asyncio.wait_for(self._process.wait(), timeout=5) - except TimeoutError: - self._process.kill() - await self._process.wait() - logger.info("Process terminated") - - @property - def is_running(self) -> bool: - return self._process is not None and self._process.returncode is None - - -class DockerProcess(LiveProcess): - """Live stdin/stdout via `docker compose exec -i`.""" - - def __init__( - self, - project_name: str, - project_dir: str, - compose_files: list[str], - service: str = "main", - ): - self._project_name = project_name - self._project_dir = project_dir - self._compose_files = compose_files - self._service = service - - @classmethod - def from_sandbox_env(cls, env: Any, service: str = "main") -> "DockerProcess": - """Create from a sandbox environment (DockerSandbox). - - ``service`` selects which compose service the agent process runs - in. Defaults to ``"main"``. Multi-container (vulhub-style) tasks - may run the agent in a dedicated attacker container — e.g. a Kali - service — while keeping target containers separate (#248). - """ - project_name = env.session_id.lower().replace(".", "-") - project_dir = str(env.environment_dir.resolve().absolute()) - compose_files = [str(p.resolve().absolute()) for p in env._docker_compose_paths] - return cls( - project_name=project_name, - project_dir=project_dir, - compose_files=compose_files, - service=service, - ) - - def _compose_cmd(self) -> list[str]: - """Base docker compose command with project/file flags.""" - cmd = [ - "docker", - "compose", - "-p", - self._project_name, - "--project-directory", - self._project_dir, - ] - for f in self._compose_files: - cmd.extend(["-f", f]) - return cmd - - def _host_env(self) -> dict[str, str]: - """Host process env with DOCKER_HOST resolved if needed.""" - proc_env = os.environ.copy() - if not proc_env.get("DOCKER_HOST"): - import subprocess - - try: - r = subprocess.run( - [ - "docker", - "context", - "inspect", - "--format", - "{{.Endpoints.docker.Host}}", - ], - capture_output=True, - text=True, - timeout=5, - ) - if r.returncode == 0 and r.stdout.strip(): - proc_env["DOCKER_HOST"] = r.stdout.strip() - except Exception: - logger.debug("Could not inspect docker context", exc_info=True) - return proc_env - - _ENV_PATH = "/tmp/.benchflow_env" - - async def _write_env_to_container( - self, - env: dict[str, str], - proc_env: dict[str, str], - ) -> None: - """Write env vars to a file inside the container (not visible in ps aux).""" - lines = "".join(f"export {k}={shlex.quote(v)}\n" for k, v in env.items()) - write_cmd = self._compose_cmd() - write_cmd.extend( - [ - "exec", - "-T", - self._service, - "bash", - "-c", - f"cat > {self._ENV_PATH} && chmod 600 {self._ENV_PATH}", - ] - ) - proc = await asyncio.create_subprocess_exec( - *write_cmd, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=proc_env, - ) - _, stderr = await asyncio.wait_for(proc.communicate(lines.encode()), timeout=30) - if proc.returncode != 0: - raise RuntimeError( - f"Failed to write env file in container (rc={proc.returncode}): " - f"{stderr.decode()[:500]}" - ) - logger.debug("Env file written inside container (%d vars)", len(env)) - - async def start( - self, - command: str, - env: dict[str, str] | None = None, - cwd: str | None = None, - ) -> None: - proc_env = self._host_env() - - # Write env vars to a file inside the container, then source it - # in the main command. This keeps secrets off `ps aux` on the host - # and avoids `--env-file` (not supported in all Compose versions). - if env: - await self._write_env_to_container(env, proc_env) - command = f"source {self._ENV_PATH} && rm -f {self._ENV_PATH} && {command}" - - cmd = self._compose_cmd() - cmd.extend(["exec", "-i", "-T"]) - if cwd: - cmd.extend(["-w", cwd]) - cmd.extend([self._service, "bash", "-c", command]) - - logger.debug(f"DockerProcess: {' '.join(cmd[:10])}...") - self._process = await asyncio.create_subprocess_exec( - *cmd, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=proc_env, - limit=_BUFFER_LIMIT, - ) - logger.info( - f"Docker process started (pid={self._process.pid}, " - f"project={self._project_name})" - ) - - -class AppleContainerProcess(LiveProcess): - """Live stdin/stdout through Apple Container's native exec transport.""" - - def __init__(self, container_name: str): - self._container_name = container_name - self._env_path = f"/tmp/.benchflow_agent_env_{uuid.uuid4().hex[:16]}" - - @classmethod - def from_sandbox_env(cls, env: Any) -> "AppleContainerProcess": - """Create from a started AppleContainerSandbox.""" - - container_name = getattr(env, "_container_name", None) - if not isinstance(container_name, str) or not container_name: - raise RuntimeError("Apple Container sandbox not started") - return cls(container_name) - - async def _write_env_to_container(self, env: dict[str, str]) -> None: - invalid = [key for key in env if not _ENV_KEY_RE.match(key)] - if invalid: - raise ValueError( - "Invalid environment variable name(s): " + ", ".join(sorted(invalid)) - ) - - lines = "".join( - f"export {key}={shlex.quote(value)}\n" for key, value in env.items() - ) - env_path = shlex.quote(self._env_path) - proc = await asyncio.create_subprocess_exec( - "container", - "exec", - "--interactive", - "--user", - "root", - self._container_name, - "sh", - "-c", - f"cat > {env_path} && chmod 600 {env_path}", - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - try: - _, stderr = await asyncio.wait_for( - proc.communicate(lines.encode()), timeout=30 - ) - except TimeoutError: - proc.kill() - await proc.wait() - raise - if proc.returncode != 0: - raise RuntimeError( - "Failed to write agent env in Apple container " - f"(rc={proc.returncode}): {stderr.decode(errors='replace')[:500]}" - ) - - async def start( - self, - command: str, - env: dict[str, str] | None = None, - cwd: str | None = None, - ) -> None: - if env: - await self._write_env_to_container(env) - env_path = shlex.quote(self._env_path) - command = f". {env_path} && rm -f {env_path} && {command}" - - args = ["container", "exec", "--interactive"] - if cwd: - args.extend(["--workdir", cwd]) - args.extend([self._container_name, "bash", "-c", command]) - self._process = await asyncio.create_subprocess_exec( - *args, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - limit=_BUFFER_LIMIT, - ) - logger.info( - "Apple Container process started (pid=%s, container=%s)", - self._process.pid, - self._container_name, - ) - - -class DaytonaProcess(LiveProcess): +class DaytonaProcess(SubprocessLiveProcess): """Live stdin/stdout via SSH to a Daytona sandbox. For DinD (compose) sandboxes, the SSH connects to the VM and then @@ -549,7 +192,7 @@ async def _cleanup_ssh_config_after_exit( self._unlink_ssh_config(path) @classmethod - async def from_sandbox_env(cls, env: Any) -> "DaytonaProcess": + async def from_sandbox_env(cls, env: Any) -> DaytonaProcess: """Create from a sandbox environment (DaytonaSandbox).""" sandbox = env._sandbox if not sandbox: @@ -761,7 +404,6 @@ class DaytonaPtyProcess(LiveProcess): direct sandboxes run the agent command directly in the PTY shell. """ - _process = None # Not used — override readline/writeline/close _START_MARKER_TIMEOUT_SEC = 120 def __init__(self, sandbox: Any, compose_cmd_prefix: str, compose_cmd_base: str): @@ -776,7 +418,7 @@ def __init__(self, sandbox: Any, compose_cmd_prefix: str, compose_cmd_base: str) self._remote_script_path: str | None = None @classmethod - async def from_sandbox_env(cls, env: Any) -> "DaytonaPtyProcess": + async def from_sandbox_env(cls, env: Any) -> DaytonaPtyProcess: sandbox = env._sandbox if not sandbox: raise RuntimeError("Daytona sandbox not started") diff --git a/src/benchflow/sandbox/process/docker.py b/src/benchflow/sandbox/process/docker.py new file mode 100644 index 000000000..1d98686e2 --- /dev/null +++ b/src/benchflow/sandbox/process/docker.py @@ -0,0 +1,157 @@ +"""Live stdio to a container managed by local Docker Compose.""" + +from __future__ import annotations + +import asyncio +import logging +import os +import shlex +from typing import Any + +from benchflow.sandbox.process._base import _BUFFER_LIMIT, SubprocessLiveProcess + +logger = logging.getLogger(__name__) + + +class DockerProcess(SubprocessLiveProcess): + """Live stdin/stdout via `docker compose exec -i`.""" + + def __init__( + self, + project_name: str, + project_dir: str, + compose_files: list[str], + service: str = "main", + ): + self._project_name = project_name + self._project_dir = project_dir + self._compose_files = compose_files + self._service = service + + @classmethod + def from_sandbox_env(cls, env: Any, service: str = "main") -> DockerProcess: + """Create from a sandbox environment (DockerSandbox). + + ``service`` selects which compose service the agent process runs + in. Defaults to ``"main"``. Multi-container (vulhub-style) tasks + may run the agent in a dedicated attacker container — e.g. a Kali + service — while keeping target containers separate (#248). + """ + project_name = env.session_id.lower().replace(".", "-") + project_dir = str(env.environment_dir.resolve().absolute()) + compose_files = [str(p.resolve().absolute()) for p in env._docker_compose_paths] + return cls( + project_name=project_name, + project_dir=project_dir, + compose_files=compose_files, + service=service, + ) + + def _compose_cmd(self) -> list[str]: + """Base docker compose command with project/file flags.""" + cmd = [ + "docker", + "compose", + "-p", + self._project_name, + "--project-directory", + self._project_dir, + ] + for f in self._compose_files: + cmd.extend(["-f", f]) + return cmd + + def _host_env(self) -> dict[str, str]: + """Host process env with DOCKER_HOST resolved if needed.""" + proc_env = os.environ.copy() + if not proc_env.get("DOCKER_HOST"): + import subprocess + + try: + r = subprocess.run( + [ + "docker", + "context", + "inspect", + "--format", + "{{.Endpoints.docker.Host}}", + ], + capture_output=True, + text=True, + timeout=5, + ) + if r.returncode == 0 and r.stdout.strip(): + proc_env["DOCKER_HOST"] = r.stdout.strip() + except Exception: + logger.debug("Could not inspect docker context", exc_info=True) + return proc_env + + _ENV_PATH = "/tmp/.benchflow_env" + + async def _write_env_to_container( + self, + env: dict[str, str], + proc_env: dict[str, str], + ) -> None: + """Write env vars to a file inside the container (not visible in ps aux).""" + lines = "".join(f"export {k}={shlex.quote(v)}\n" for k, v in env.items()) + write_cmd = self._compose_cmd() + write_cmd.extend( + [ + "exec", + "-T", + self._service, + "bash", + "-c", + f"cat > {self._ENV_PATH} && chmod 600 {self._ENV_PATH}", + ] + ) + proc = await asyncio.create_subprocess_exec( + *write_cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=proc_env, + ) + _, stderr = await asyncio.wait_for(proc.communicate(lines.encode()), timeout=30) + if proc.returncode != 0: + raise RuntimeError( + f"Failed to write env file in container (rc={proc.returncode}): " + f"{stderr.decode()[:500]}" + ) + logger.debug("Env file written inside container (%d vars)", len(env)) + + async def start( + self, + command: str, + env: dict[str, str] | None = None, + cwd: str | None = None, + ) -> None: + proc_env = self._host_env() + + # Write env vars to a file inside the container, then source it + # in the main command. This keeps secrets off `ps aux` on the host + # and avoids `--env-file` (not supported in all Compose versions). + if env: + await self._write_env_to_container(env, proc_env) + command = f"source {self._ENV_PATH} && rm -f {self._ENV_PATH} && {command}" + + cmd = self._compose_cmd() + cmd.extend(["exec", "-i", "-T"]) + if cwd: + cmd.extend(["-w", cwd]) + cmd.extend([self._service, "bash", "-c", command]) + + logger.debug(f"DockerProcess: {' '.join(cmd[:10])}...") + self._process = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=proc_env, + limit=_BUFFER_LIMIT, + ) + logger.info( + f"Docker process started (pid={self._process.pid}, " + f"project={self._project_name})" + ) diff --git a/src/benchflow/sandbox/providers.py b/src/benchflow/sandbox/providers.py index ed90c21fc..1ee2e34f8 100644 --- a/src/benchflow/sandbox/providers.py +++ b/src/benchflow/sandbox/providers.py @@ -36,6 +36,11 @@ class SandboxProvider: name: str extra: str | None # pip/uv optional-dependency extra; None when none is needed model_proxy: ModelProxyLocation + #: Whether the backend can actually enforce ``network_mode = "no-network"``. + #: Declared here so the runtime capability gate reads a fact off the + #: registry instead of growing a ``sandbox == ""`` special case per + #: backend that cannot isolate the network. + enforces_no_network: bool = True @property def off_box_model(self) -> bool: @@ -61,6 +66,16 @@ def off_box_model(self) -> bool: "apple-container", extra=None, model_proxy=ModelProxyLocation.SANDBOX, + enforces_no_network=False, + ), + SandboxProvider( + "agentcore", + extra="sandbox-agentcore", + model_proxy=ModelProxyLocation.SANDBOX, + # AgentCore's CreateAgentRuntime networkConfiguration.networkMode only + # accepts PUBLIC or VPC — there is no isolated mode, so a no-network + # task cannot be honored here. + enforces_no_network=False, ), ) @@ -82,6 +97,12 @@ def off_box_model(self) -> bool: OFF_BOX_MODEL_PROVIDERS = SANDBOX_MODEL_PROXY_PROVIDERS +#: Providers that cannot enforce ``network_mode = "no-network"``. +NO_NETWORK_UNSUPPORTED_PROVIDERS: frozenset[str] = frozenset( + p.name for p in _PROVIDERS if not p.enforces_no_network +) + + def is_known_provider(name: str) -> bool: """True if ``name`` is a registered sandbox provider.""" return name in SANDBOX_PROVIDER_SET diff --git a/src/benchflow/sandbox/setup.py b/src/benchflow/sandbox/setup.py index 1f5bdc28a..746399cea 100644 --- a/src/benchflow/sandbox/setup.py +++ b/src/benchflow/sandbox/setup.py @@ -781,6 +781,23 @@ def _create_sandbox_environment( task_env_config=env_config, persistent_env=manifest_env or None, ) + elif sandbox_type == "agentcore": + try: + import bedrock_agentcore # noqa: F401 + except ImportError as exc: + _raise_missing_optional_sandbox_dependency("agentcore", exc) + + from benchflow.sandbox.agentcore import AgentCoreSandbox + + AgentCoreSandbox.preflight() + return AgentCoreSandbox( + environment_dir=environment_dir, + environment_name=task_path.name, + session_id=rollout_name, + rollout_paths=rollout_paths, + task_env_config=env_config, + persistent_env=manifest_env or None, + ) elif sandbox_type == "apple-container": from benchflow.sandbox.apple_container import AppleContainerSandbox diff --git a/src/benchflow/task/runtime_capabilities.py b/src/benchflow/task/runtime_capabilities.py index 2afd0c114..ba1f74d65 100644 --- a/src/benchflow/task/runtime_capabilities.py +++ b/src/benchflow/task/runtime_capabilities.py @@ -15,7 +15,11 @@ from typing import cast from benchflow.rewards.rubric_config import criteria_aggregate_policy_from_rubric -from benchflow.sandbox.providers import SANDBOX_PROVIDER_SET, providers_phrase +from benchflow.sandbox.providers import ( + NO_NETWORK_UNSUPPORTED_PROVIDERS, + SANDBOX_PROVIDER_SET, + providers_phrase, +) from benchflow.task.config import ( NetworkMode, TaskConfig, @@ -261,11 +265,11 @@ def _append_network_issue( mode: NetworkMode | None, sandbox: str, ) -> None: - if sandbox == "apple-container" and mode == NetworkMode.NO_NETWORK: + if mode == NetworkMode.NO_NETWORK and sandbox in NO_NETWORK_UNSUPPORTED_PROVIDERS: _issue( unsupported, path=path, - reason="network_mode='no-network' is not enforced by apple-container", + reason=f"network_mode='no-network' is not enforced by {sandbox}", sandbox=sandbox, ) if mode == NetworkMode.ALLOWLIST: diff --git a/tests/test_acp.py b/tests/test_acp.py index 97f51790c..f6b9c68a6 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -957,10 +957,6 @@ async def test_model_id_selection(self, model_in, expected_model, tmp_path): mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -986,10 +982,6 @@ async def test_pi_acp_preserves_registered_provider_prefix(self, tmp_path): mock_acp = self._make_mocks() mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -1015,10 +1007,6 @@ async def test_opencode_keeps_modelsdev_formatting(self, tmp_path): mock_acp = self._make_mocks() mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -1043,10 +1031,6 @@ async def test_openhands_skips_set_model(self, tmp_path): mock_acp = self._make_mocks() mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -1076,10 +1060,6 @@ async def test_openhands_direct_execution_patches_before_privilege_drop( mock_env.exec.return_value = MagicMock(return_code=0, stdout="", stderr="") with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch( "benchflow.acp.runtime.ContainerTransport", return_value=MagicMock(), @@ -1122,10 +1102,6 @@ async def test_codex_uses_session_advertised_model_id(self, tmp_path): } mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -1155,10 +1131,6 @@ async def test_claude_uses_config_options_for_model_and_effort(self, tmp_path): ] mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -1189,10 +1161,6 @@ async def test_claude_litellm_env_owns_model_selection(self, tmp_path): mock_acp.session_new.return_value.config_options = [{"id": "model"}] mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -1217,25 +1185,24 @@ async def test_claude_litellm_env_owns_model_selection(self, tmp_path): mock_acp.set_config_option.assert_not_awaited() @pytest.mark.asyncio - async def test_apple_container_uses_native_live_process(self, tmp_path): - """Guards PR #936 against treating Apple Container as Daytona.""" - + async def test_connect_acp_delegates_transport_to_the_sandbox(self, tmp_path): + """connect_acp asks the sandbox for its transport and wires it through. + + Per-backend transport selection moved onto the sandbox objects; the + cases that used to be asserted here (Apple Container vs Daytona, + PTY vs SSH, the gemini and BENCHFLOW_DAYTONA_ACP_TRANSPORT overrides) + now live in ``tests/test_sandbox_live_process.py``. What this layer + still owns is the delegation itself. + """ from benchflow.acp.runtime import connect_acp mock_acp = self._make_mocks() + live_process = MagicMock() mock_env = MagicMock() mock_env.exec = AsyncMock(return_value=MagicMock(return_code=1, stdout="")) - live_process = MagicMock() + mock_env.live_process = AsyncMock(return_value=live_process) with ( - patch( - "benchflow.acp.runtime.AppleContainerProcess.from_sandbox_env", - return_value=live_process, - ) as apple_process, - patch( - "benchflow.acp.runtime.DaytonaProcess.from_sandbox_env", - new_callable=AsyncMock, - ) as daytona_process, patch("benchflow.acp.runtime.ContainerTransport") as transport, patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -1251,208 +1218,9 @@ async def test_apple_container_uses_native_live_process(self, tmp_path): agent_cwd="/root", ) - apple_process.assert_called_once_with(mock_env) - daytona_process.assert_not_awaited() + mock_env.live_process.assert_awaited_once_with(agent="codex-acp") assert transport.call_args.kwargs["container_process"] is live_process - @pytest.mark.asyncio - async def test_daytona_dind_uses_pty_transport(self, tmp_path): - """Daytona compose tasks use PTY transport to avoid SSH pipe-closed failures.""" - from benchflow.acp.runtime import connect_acp - - mock_acp = self._make_mocks() - mock_env = MagicMock() - mock_env.exec = AsyncMock(return_value=MagicMock(return_code=1, stdout="")) - mock_env._strategy = MagicMock() - mock_env._strategy._compose_cmd = MagicMock(return_value="docker compose -p t") - - with ( - patch( - "benchflow.acp.runtime.DaytonaPtyProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_pty, - patch( - "benchflow.acp.runtime.DaytonaProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_ssh, - patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), - patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), - ): - await connect_acp( - env=mock_env, - agent="test-agent", - agent_launch="test-agent", - agent_env={}, - sandbox_user=None, - model=None, - rollout_dir=tmp_path, - environment="daytona", - agent_cwd="/app", - ) - - mock_pty.assert_awaited_once_with(mock_env) - mock_ssh.assert_not_awaited() - - @pytest.mark.asyncio - async def test_daytona_direct_uses_pty_transport(self, tmp_path): - """Direct Daytona tasks also use PTY transport, not SSH pipes.""" - from benchflow.acp.runtime import connect_acp - - mock_acp = self._make_mocks() - mock_env = MagicMock() - mock_env.exec = AsyncMock(return_value=MagicMock(return_code=1, stdout="")) - - with ( - patch( - "benchflow.acp.runtime.DaytonaPtyProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_pty, - patch( - "benchflow.acp.runtime.DaytonaProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_ssh, - patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), - patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), - ): - await connect_acp( - env=mock_env, - agent="test-agent", - agent_launch="test-agent", - agent_env={}, - sandbox_user=None, - model=None, - rollout_dir=tmp_path, - environment="daytona", - agent_cwd="/app", - ) - - mock_pty.assert_awaited_once_with(mock_env) - mock_ssh.assert_not_awaited() - - @pytest.mark.asyncio - async def test_daytona_direct_can_opt_into_ssh_transport( - self, tmp_path, monkeypatch - ): - """Guards PR #921 fallback for PTY post-tool controller deadlocks.""" - from benchflow.acp.runtime import connect_acp - - monkeypatch.setenv("BENCHFLOW_DAYTONA_ACP_TRANSPORT", "ssh") - mock_acp = self._make_mocks() - mock_env = MagicMock() - mock_env.exec = AsyncMock(return_value=MagicMock(return_code=1, stdout="")) - - with ( - patch( - "benchflow.acp.runtime.DaytonaPtyProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_pty, - patch( - "benchflow.acp.runtime.DaytonaProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_ssh, - patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), - patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), - ): - await connect_acp( - env=mock_env, - agent="openhands", - agent_launch="openhands acp", - agent_env={}, - sandbox_user=None, - model=None, - rollout_dir=tmp_path, - environment="daytona", - agent_cwd="/app", - ) - - mock_ssh.assert_awaited_once_with(mock_env) - mock_pty.assert_not_awaited() - - @pytest.mark.asyncio - async def test_invalid_daytona_transport_falls_back_to_pty( - self, tmp_path, monkeypatch - ): - """Guards PR #921 against invalid transport config disabling Daytona.""" - from benchflow.acp.runtime import connect_acp - - monkeypatch.setenv("BENCHFLOW_DAYTONA_ACP_TRANSPORT", "invalid") - mock_acp = self._make_mocks() - mock_env = MagicMock() - mock_env.exec = AsyncMock(return_value=MagicMock(return_code=1, stdout="")) - - with ( - patch( - "benchflow.acp.runtime.DaytonaPtyProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_pty, - patch( - "benchflow.acp.runtime.DaytonaProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_ssh, - patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), - patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), - ): - await connect_acp( - env=mock_env, - agent="openhands", - agent_launch="openhands acp", - agent_env={}, - sandbox_user=None, - model=None, - rollout_dir=tmp_path, - environment="daytona", - agent_cwd="/app", - ) - - mock_pty.assert_awaited_once_with(mock_env) - mock_ssh.assert_not_awaited() - - @pytest.mark.asyncio - async def test_daytona_gemini_uses_ssh_transport(self, tmp_path): - """Guards the Gemini regression introduced by PR #896's PTY migration.""" - from benchflow.acp.runtime import connect_acp - - mock_acp = self._make_mocks() - mock_env = MagicMock() - mock_env.exec = AsyncMock(return_value=MagicMock(return_code=1, stdout="")) - - with ( - patch( - "benchflow.acp.runtime.DaytonaPtyProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_pty, - patch( - "benchflow.acp.runtime.DaytonaProcess.from_sandbox_env", - new_callable=AsyncMock, - return_value=MagicMock(), - ) as mock_ssh, - patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), - patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), - ): - await connect_acp( - env=mock_env, - agent="gemini", - agent_launch="gemini --acp --yolo", - agent_env={"GEMINI_API_KEY": "test"}, - sandbox_user=None, - model="gemini-3.1-flash-lite", - rollout_dir=tmp_path, - environment="daytona", - agent_cwd="/app", - ) - - mock_ssh.assert_awaited_once_with(mock_env) - mock_pty.assert_not_awaited() - class TestSandboxStartupDiagnostics: """Guards ENG-147: sandbox startup failures must carry structured diagnostics.""" @@ -1586,7 +1354,7 @@ async def test_live_process_raises_typed_transport_error_with_rc(self) -> None: """``LiveProcess.readline`` raises ``TransportClosedError`` with the structured ``process_exit_code`` filled in (issue #504).""" from benchflow.diagnostics import TransportClosedError - from benchflow.sandbox.process import LiveProcess + from benchflow.sandbox.process import SubprocessLiveProcess class _StubProcess: def __init__(self) -> None: @@ -1601,7 +1369,7 @@ async def readline(self) -> bytes: async def read(self, _n: int) -> bytes: return b"Connection to sandbox lost" - class _LP(LiveProcess): + class _LP(SubprocessLiveProcess): async def start( self, command, env=None, cwd=None ) -> None: # pragma: no cover - unused @@ -1624,7 +1392,7 @@ async def test_live_process_raises_typed_transport_error_when_remote_killed( """rc=None ⇒ remote session was killed; the diagnostic must carry the pid and the ``remote_session_killed`` diagnosis (issue #504).""" from benchflow.diagnostics import TransportClosedError - from benchflow.sandbox.process import LiveProcess + from benchflow.sandbox.process import SubprocessLiveProcess class _StubProcess: def __init__(self) -> None: @@ -1636,7 +1404,7 @@ def __init__(self) -> None: async def readline(self) -> bytes: return b"" - class _LP(LiveProcess): + class _LP(SubprocessLiveProcess): async def start( self, command, env=None, cwd=None ) -> None: # pragma: no cover - unused diff --git a/tests/test_acp_model_config_dispatch.py b/tests/test_acp_model_config_dispatch.py index a05fdfc5a..88663ea13 100644 --- a/tests/test_acp_model_config_dispatch.py +++ b/tests/test_acp_model_config_dispatch.py @@ -43,10 +43,6 @@ def _make_mocks(config_options=None, model_state=None): @contextlib.contextmanager def _runtime_patches(mock_acp): with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): diff --git a/tests/test_acp_setup_failure_propagation.py b/tests/test_acp_setup_failure_propagation.py index 6353cfb93..b5987d34f 100644 --- a/tests/test_acp_setup_failure_propagation.py +++ b/tests/test_acp_setup_failure_propagation.py @@ -47,10 +47,6 @@ async def test_set_model_failure_aborts_rollout(tmp_path) -> None: mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), pytest.raises(RuntimeError, match="Failed to set model"), @@ -82,10 +78,6 @@ async def test_config_option_failure_aborts_rollout(tmp_path) -> None: mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), pytest.raises(RuntimeError, match="Failed to set ACP model config option"), @@ -116,10 +108,6 @@ async def test_set_model_timeout_aborts_rollout(tmp_path) -> None: mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), pytest.raises(RuntimeError, match="Failed to set model"), @@ -148,10 +136,6 @@ async def test_set_model_success_still_returns_session(tmp_path) -> None: mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -186,10 +170,6 @@ async def test_no_model_does_not_call_set_model(tmp_path) -> None: mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), ): @@ -220,10 +200,6 @@ async def test_initialize_timeout_is_transport_failure_not_agent_timeout( mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), patch("benchflow.acp.runtime.asyncio.sleep", new_callable=AsyncMock), @@ -264,10 +240,6 @@ async def enforce(*args, **kwargs): mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch("benchflow.acp.runtime.ContainerTransport", return_value=MagicMock()), patch("benchflow.acp.runtime.ACPClient", return_value=mock_acp), patch( diff --git a/tests/test_agentcore_sandbox.py b/tests/test_agentcore_sandbox.py new file mode 100644 index 000000000..d4a6edc9c --- /dev/null +++ b/tests/test_agentcore_sandbox.py @@ -0,0 +1,419 @@ +"""AgentCore sandbox behaviour, driven against a fake boto3 command plane. + +The expectations encoded here were measured against the live service during +the design of this backend (issue #935), not inferred from documentation: + +* ``contentDelta`` carries ``stdout`` and ``stderr`` as separate fields, so + ``ExecResult`` must keep them separate. +* ``contentStop.exitCode`` is the command's real exit status and must not be + flattened to 0. +* ``contentStop.status == "TIMED_OUT"`` is how a timeout is reported. +* Throttling and quota exhaustion arrive as typed members of the response + stream, and are infrastructure failures rather than failed commands. + +The live end-to-end path is exercised separately by +``test_real_agentcore_lifecycle``, which is skipped unless +``BENCHFLOW_AGENTCORE_LIVE_TEST=1`` and AWS credentials are present. +""" + +from __future__ import annotations + +import base64 +import os +import tarfile +from unittest.mock import MagicMock, patch + +import pytest + +from benchflow.sandbox.agentcore import AgentCoreSandbox +from benchflow.sandbox.protocol import SandboxStartupError +from benchflow.task.config import SandboxConfig + + +def _stream(*, stdout="", stderr="", exit_code=0, status="COMPLETED"): + chunks = [] + if stdout or stderr: + chunks.append({"chunk": {"contentDelta": {"stdout": stdout, "stderr": stderr}}}) + chunks.append({"chunk": {"contentStop": {"exitCode": exit_code, "status": status}}}) + return {"stream": chunks} + + +@pytest.fixture +def sandbox(tmp_path): + (tmp_path / "Dockerfile").write_text("FROM python:3.12-slim\n") + env = AgentCoreSandbox( + environment_dir=tmp_path, + environment_name="demo-task", + session_id="run-1", + rollout_paths=None, + task_env_config=SandboxConfig(), + ) + env.runtime_arn = "arn:aws:bedrock-agentcore:us-west-2:1:runtime/demo" + env.runtime_session_id = "s" * 40 + return env + + +class TestExecSemantics: + @pytest.mark.asyncio + async def test_stdout_and_stderr_stay_separate(self, sandbox): + """The service splits the streams; ExecResult must not merge them.""" + client = MagicMock() + client.invoke_agent_runtime_command.return_value = _stream( + stdout="out", stderr="err", exit_code=0 + ) + with patch.object(sandbox, "_client", return_value=client): + result = await sandbox.exec("echo hi") + + assert result.stdout == "out" + assert result.stderr == "err" + assert result.return_code == 0 + + @pytest.mark.asyncio + async def test_nonzero_exit_code_is_preserved(self, sandbox): + client = MagicMock() + client.invoke_agent_runtime_command.return_value = _stream(exit_code=7) + with patch.object(sandbox, "_client", return_value=client): + result = await sandbox.exec("exit 7") + + assert result.return_code == 7 + + @pytest.mark.asyncio + async def test_missing_content_stop_is_not_reported_as_success(self, sandbox): + """A truncated stream has no exit status — inventing 0 would hide it.""" + client = MagicMock() + client.invoke_agent_runtime_command.return_value = { + "stream": [{"chunk": {"contentDelta": {"stdout": "partial"}}}] + } + with patch.object(sandbox, "_client", return_value=client): + result = await sandbox.exec("something") + + assert result.return_code != 0 + + @pytest.mark.asyncio + async def test_timed_out_status_raises(self, sandbox): + client = MagicMock() + client.invoke_agent_runtime_command.return_value = _stream( + exit_code=None, status="TIMED_OUT" + ) + with ( + patch.object(sandbox, "_client", return_value=client), + pytest.raises(TimeoutError), + ): + await sandbox.exec("sleep 999", timeout_sec=5) + + @pytest.mark.asyncio + async def test_command_is_wrapped_for_bash(self, sandbox): + """AgentCore does not run a shell for you; commands need bash -c.""" + client = MagicMock() + client.invoke_agent_runtime_command.return_value = _stream() + with patch.object(sandbox, "_client", return_value=client): + await sandbox.exec("echo hi") + + body = client.invoke_agent_runtime_command.call_args.kwargs["body"] + assert body["command"].startswith("/bin/bash -c ") + + @pytest.mark.asyncio + async def test_timeout_is_clamped_to_service_range(self, sandbox): + """The service rejects timeouts outside 1..3600.""" + client = MagicMock() + client.invoke_agent_runtime_command.return_value = _stream() + with patch.object(sandbox, "_client", return_value=client): + await sandbox.exec("echo hi", timeout_sec=99999) + + body = client.invoke_agent_runtime_command.call_args.kwargs["body"] + assert body["timeout"] == 3600 + + @pytest.mark.asyncio + async def test_secrets_are_not_passed_as_literal_argv(self, sandbox): + """Env values must go through the base64 env-file wrapper (#412).""" + client = MagicMock() + client.invoke_agent_runtime_command.return_value = _stream() + with patch.object(sandbox, "_client", return_value=client): + await sandbox.exec("run", env={"SECRET_TOKEN": "hunter2"}) + + command = client.invoke_agent_runtime_command.call_args.kwargs["body"][ + "command" + ] + assert "hunter2" not in command + assert "base64 -d" in command + + @pytest.mark.asyncio + async def test_non_main_service_is_rejected(self, sandbox): + with pytest.raises(ValueError, match="single-container"): + await sandbox.exec("echo hi", service="target") + + +class TestInfrastructureAttribution: + @pytest.mark.parametrize( + "error_key", + ["throttlingException", "serviceQuotaExceededException"], + ) + @pytest.mark.asyncio + async def test_capacity_errors_are_infra_not_agent_failure( + self, sandbox, error_key + ): + """Throttling must not be recorded as a command that failed. + + Attributing a quota error to the agent is how an infrastructure + problem silently becomes a scored 0 in result.json. + """ + client = MagicMock() + client.invoke_agent_runtime_command.return_value = { + "stream": [{error_key: {"message": "slow down"}}] + } + with ( + patch.object(sandbox, "_client", return_value=client), + pytest.raises(SandboxStartupError), + ): + await sandbox.exec("echo hi") + + +class TestSessionWarmup: + @pytest.mark.asyncio + async def test_cold_start_500_is_retried_not_fatal(self, sandbox, monkeypatch): + """A cold image 500s on the first command; that is not a task failure. + + READY describes the runtime definition, not a running session — the + first command pulls the image and starts the container. Observed live: + a never-pulled image fails the first command and succeeds moments + later. Failing closed here would score a rollout 0 for a sandbox that + was merely still booting. + """ + monkeypatch.setattr( + "benchflow.sandbox.agentcore._SESSION_WARMUP_BACKOFF_SEC", 0.0 + ) + calls = {"n": 0} + + async def _exec(command, **kwargs): + calls["n"] += 1 + if calls["n"] < 3: + raise RuntimeError("Received error (500) from runtime") + return MagicMock(return_code=0, stdout="", stderr="") + + with patch.object(sandbox, "exec", side_effect=_exec): + await sandbox._warm_session() + + assert calls["n"] == 3 + + @pytest.mark.asyncio + async def test_persistent_failure_still_reports_startup_error( + self, sandbox, monkeypatch + ): + monkeypatch.setattr( + "benchflow.sandbox.agentcore._SESSION_WARMUP_BACKOFF_SEC", 0.0 + ) + + async def _exec(command, **kwargs): + raise RuntimeError("Received error (500) from runtime") + + with ( + patch.object(sandbox, "exec", side_effect=_exec), + pytest.raises(SandboxStartupError, match="did not accept commands"), + ): + await sandbox._warm_session() + + @pytest.mark.asyncio + async def test_throttling_during_warmup_is_not_retried(self, sandbox): + """Quota errors are infra and terminal; retrying only wastes time.""" + + async def _exec(command, **kwargs): + raise SandboxStartupError("AgentCore throttlingException: slow down") + + with ( + patch.object(sandbox, "exec", side_effect=_exec) as mock_exec, + pytest.raises(SandboxStartupError, match="throttling"), + ): + await sandbox._warm_session() + + assert mock_exec.call_count == 1 + + +class TestFileTransfer: + @pytest.mark.asyncio + async def test_upload_dir_ships_one_archive_not_one_command_per_file( + self, sandbox, tmp_path + ): + """A round trip per file would be unusably slow for a task tree.""" + source = tmp_path / "payload" + (source / "nested").mkdir(parents=True) + for i in range(12): + (source / f"file{i}.txt").write_text(f"contents {i}") + (source / "nested" / "deep.txt").write_text("deep") + + commands: list[str] = [] + + async def _record(command, **kwargs): + commands.append(command) + return MagicMock(return_code=0, stdout="", stderr="") + + with patch.object(sandbox, "exec", side_effect=_record): + await sandbox.upload_dir(source, "/workspace") + + # mkdir + staged chunk(s) + one extract, not 13 uploads. + assert len(commands) < 13 + assert any("tar -xzf" in c for c in commands) + + @pytest.mark.asyncio + async def test_upload_dir_skips_symlinks(self, sandbox, tmp_path): + """Guards #411: a task symlink must not exfiltrate host files.""" + source = tmp_path / "payload" + source.mkdir() + (source / "real.txt").write_text("ok") + secret = tmp_path / "host-secret.txt" + secret.write_text("do not ship me") + (source / "link.txt").symlink_to(secret) + + staged: list[str] = [] + + async def _record(command, **kwargs): + staged.append(command) + return MagicMock(return_code=0, stdout="", stderr="") + + with patch.object(sandbox, "exec", side_effect=_record): + await sandbox.upload_dir(source, "/workspace") + + blob = "".join(staged) + assert "do not ship me" not in blob + + @pytest.mark.asyncio + async def test_download_dir_extracts_the_returned_archive(self, sandbox, tmp_path): + payload = tmp_path / "src" + payload.mkdir() + (payload / "result.json").write_text('{"reward": 1.0}') + archive = tmp_path / "a.tar.gz" + with tarfile.open(archive, "w:gz") as tar: + tar.add(payload / "result.json", arcname="result.json") + encoded = base64.b64encode(archive.read_bytes()).decode() + + async def _exec(command, **kwargs): + return MagicMock(return_code=0, stdout=encoded, stderr="") + + target = tmp_path / "out" + with patch.object(sandbox, "exec", side_effect=_exec): + await sandbox.download_dir("/logs", target) + + assert (target / "result.json").read_text() == '{"reward": 1.0}' + + @pytest.mark.asyncio + async def test_oversized_inline_write_is_refused(self, sandbox): + """Silently truncating past the 64 KB command cap would corrupt files.""" + with pytest.raises(ValueError, match="Refusing to inline"): + await sandbox.write_text_file("/tmp/big", "x" * 200_000) + + +class TestImagePreparation: + def test_generated_dockerfile_adds_the_ping_shim(self, sandbox, tmp_path): + """Without a /ping responder the microVM 500s on every command.""" + context, dockerfile, shim = sandbox._materialize_build_context() + text = dockerfile.read_text() + + assert "FROM python:3.12-slim" in text + assert "benchflow_agentcore_shim.py" in text + assert "EXPOSE 8080" in text + assert shim.exists() + assert shim.parent == context + + def test_task_dockerfile_is_not_modified_in_place(self, sandbox, tmp_path): + original = (tmp_path / "Dockerfile").read_text() + sandbox._materialize_build_context() + + assert (tmp_path / "Dockerfile").read_text() == original + + @pytest.mark.asyncio + async def test_build_leaves_no_generated_files_in_the_task_tree( + self, sandbox, tmp_path + ): + """The build context is the task's own directory; don't litter it.""" + before = {p.name for p in tmp_path.iterdir()} + + async def _to_thread(fn, *args, **kwargs): + return MagicMock(returncode=0, stdout="sha256:" + "a" * 64, stderr="") + + with ( + patch("benchflow.sandbox.agentcore.asyncio.to_thread", new=_to_thread), + patch.object(sandbox, "_ecr_registry", return_value="reg"), + ): + await sandbox._publish_image(force_build=False) + + assert {p.name for p in tmp_path.iterdir()} == before + + def test_docker_image_config_is_honoured(self, tmp_path): + env = AgentCoreSandbox( + environment_dir=tmp_path, + environment_name="img-task", + session_id="run-1", + rollout_paths=None, + task_env_config=SandboxConfig(docker_image="python:3.12-slim"), + ) + _context, dockerfile, _shim = env._materialize_build_context() + + assert "FROM python:3.12-slim" in dockerfile.read_text() + + +class TestCapabilityGating: + def test_snapshots_are_unsupported(self, sandbox): + """AgentCore has no container checkpoint primitive.""" + assert sandbox.supports_snapshot is False + + def test_no_network_tasks_are_refused(self, tmp_path): + """networkMode offers only PUBLIC/VPC, so isolation cannot be honoured.""" + from benchflow.task.config import TaskConfig + from benchflow.task.runtime_capabilities import validate_task_runtime_support + + config = TaskConfig.model_validate( + {"environment": {"network_mode": "no-network"}} + ) + issues = validate_task_runtime_support(config, sandbox="agentcore") + + assert any( + "no-network' is not enforced by agentcore" in issue.reason + for issue in issues + ) + + +@pytest.mark.skipif( + os.environ.get("BENCHFLOW_AGENTCORE_LIVE_TEST") != "1", + reason="live AWS test; set BENCHFLOW_AGENTCORE_LIVE_TEST=1 to run", +) +@pytest.mark.asyncio +async def test_real_agentcore_lifecycle(tmp_path): + """End-to-end against the live service: build, run, transfer, tear down. + + Requires AWS credentials with bedrock-agentcore and ECR access, a docker + daemon, and BENCHFLOW_AGENTCORE_ROLE_ARN pointing at the runtime execution + role. + """ + (tmp_path / "Dockerfile").write_text( + "FROM python:3.12-slim\nRUN echo baked > /baked.txt\n" + ) + env = AgentCoreSandbox( + environment_dir=tmp_path, + environment_name="live-canary", + session_id="live-1", + rollout_paths=None, + task_env_config=SandboxConfig(), + ) + AgentCoreSandbox.preflight() + await env.start(force_build=False) + try: + baked = await env.exec("cat /baked.txt") + assert baked.return_code == 0 + assert "baked" in (baked.stdout or "") + + streams = await env.exec("echo o; echo e 1>&2; exit 3") + assert streams.return_code == 3 + assert "o" in (streams.stdout or "") + assert "e" in (streams.stderr or "") + + payload = tmp_path / "payload" + payload.mkdir() + (payload / "hello.txt").write_text("from-host") + await env.upload_dir(payload, "/workspace/payload") + echoed = await env.exec("cat /workspace/payload/hello.txt") + assert (echoed.stdout or "").strip() == "from-host" + + out = tmp_path / "out" + await env.download_dir("/workspace/payload", out) + assert (out / "hello.txt").read_text() == "from-host" + finally: + await env.stop(delete=True) diff --git a/tests/test_oracle_chokepoint.py b/tests/test_oracle_chokepoint.py index f15557c2b..bd4c676fe 100644 --- a/tests/test_oracle_chokepoint.py +++ b/tests/test_oracle_chokepoint.py @@ -34,6 +34,16 @@ def _plain_cli_output(output: str) -> str: return ANSI_ESCAPE_RE.sub("", output) +# Rich draws help panels with box-drawing borders and hard-wraps long option +# help across lines. Removing the borders and collapsing whitespace recovers +# the logical text so assertions can match a phrase that happens to wrap. +_BOX_DRAWING_RE = re.compile(r"[─-╿]") + + +def _unwrapped_cli_output(output: str) -> str: + return " ".join(_BOX_DRAWING_RE.sub(" ", _plain_cli_output(output)).split()) + + class TestEvalCreateRouting: """`bench eval run` must dispatch to cli/main.py:eval_run. @@ -76,7 +86,12 @@ def test_sandbox_help_matches_registered_backends(self, command): result = CliRunner().invoke(app, command) assert result.exit_code == 0 - assert f"Sandbox: {providers_phrase()}" in result.stdout + # Rich wraps the options box at the terminal width, so once enough + # providers are registered the phrase spans lines and picks up the + # box border between them. Strip the borders and collapse whitespace + # so this stays a test about the provider list rather than about how + # wide the terminal happened to be. + assert f"Sandbox: {providers_phrase()}" in _unwrapped_cli_output(result.stdout) assert "firecracker" not in result.stdout.lower() assert "kubernetes" not in result.stdout.lower() assert "k8s" not in result.stdout.lower() diff --git a/tests/test_process.py b/tests/test_process.py index 244f8d9bc..ccdbd3a09 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -360,7 +360,7 @@ async def test_native_exec_is_interactive_and_uses_workdir(self): calls = [] inputs = [] with patch( - "benchflow.sandbox.process.asyncio.create_subprocess_exec", + "benchflow.sandbox.process.apple.asyncio.create_subprocess_exec", side_effect=self._fake_exec(calls, inputs), ): proc = AppleContainerProcess("bf_run") @@ -388,7 +388,7 @@ async def test_secret_env_uses_stdin_not_process_argv(self): calls = [] inputs = [] with patch( - "benchflow.sandbox.process.asyncio.create_subprocess_exec", + "benchflow.sandbox.process.apple.asyncio.create_subprocess_exec", side_effect=self._fake_exec(calls, inputs), ): proc = AppleContainerProcess("bf_run") diff --git a/tests/test_rollout_on_ask_user_wiring.py b/tests/test_rollout_on_ask_user_wiring.py index aa18f706a..916b712f4 100644 --- a/tests/test_rollout_on_ask_user_wiring.py +++ b/tests/test_rollout_on_ask_user_wiring.py @@ -217,10 +217,6 @@ def _capture_handler(handler: Any) -> None: mock_env = AsyncMock() with ( - patch( - "benchflow.acp.runtime.DockerProcess.from_sandbox_env", - return_value=MagicMock(), - ), patch( "benchflow.acp.runtime.ContainerTransport", return_value=MagicMock(), diff --git a/tests/test_runtime_capabilities.py b/tests/test_runtime_capabilities.py index 4d1ec44f4..0b138811b 100644 --- a/tests/test_runtime_capabilities.py +++ b/tests/test_runtime_capabilities.py @@ -159,7 +159,8 @@ def test_validator_reports_unknown_sandbox_backend() -> None: assert [(issue.path, issue.reason) for issue in issues] == [ ( "sandbox", - "unknown sandbox backend; use docker, daytona, modal, or apple-container", + "unknown sandbox backend; use docker, daytona, modal, apple-container, " + "or agentcore", ) ] diff --git a/tests/test_sandbox_live_process.py b/tests/test_sandbox_live_process.py new file mode 100644 index 000000000..41665e28d --- /dev/null +++ b/tests/test_sandbox_live_process.py @@ -0,0 +1,224 @@ +"""Each sandbox backend owns its ACP transport choice. + +These cases previously lived in ``tests/test_acp.py`` and drove +``connect_acp`` with a provider-name string, because transport selection was +an ``if environment == ...`` chain inside the ACP layer. That chain is now a +single ``await env.live_process(agent=...)`` call, so the same guarantees are +asserted directly against the backend that makes the decision. The regression +each case guards is unchanged and named in its docstring. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from benchflow.sandbox import process as process_pkg +from benchflow.sandbox._base import BaseSandbox + + +def _stub_sandbox(**attrs: object) -> SimpleNamespace: + """A stand-in for a started sandbox. + + ``live_process`` only forwards ``self`` to the transport's + ``from_sandbox_env``, so the concrete sandbox never has to be built. + """ + return SimpleNamespace(**attrs) + + +class TestDaytonaTransportSelection: + @pytest.mark.asyncio + async def test_direct_uses_pty_transport(self) -> None: + """Direct Daytona tasks use PTY transport, not SSH pipes.""" + from benchflow.sandbox.daytona import DaytonaSandbox + + env = _stub_sandbox() + with ( + patch.object( + process_pkg.DaytonaPtyProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as pty, + patch.object( + process_pkg.DaytonaProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as ssh, + ): + await DaytonaSandbox.live_process(env, agent="test-agent") + + pty.assert_awaited_once_with(env) + ssh.assert_not_awaited() + + @pytest.mark.asyncio + async def test_compose_task_uses_pty_transport(self) -> None: + """Daytona compose (DinD) tasks avoid SSH pipe-closed failures.""" + from benchflow.sandbox.daytona import DaytonaSandbox + + strategy = MagicMock() + strategy._compose_cmd = MagicMock(return_value="docker compose -p t") + env = _stub_sandbox(_strategy=strategy) + with ( + patch.object( + process_pkg.DaytonaPtyProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as pty, + patch.object( + process_pkg.DaytonaProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as ssh, + ): + await DaytonaSandbox.live_process(env, agent="test-agent") + + pty.assert_awaited_once_with(env) + ssh.assert_not_awaited() + + @pytest.mark.asyncio + async def test_can_opt_into_ssh_transport(self, monkeypatch) -> None: + """Guards PR #921 fallback for PTY post-tool controller deadlocks.""" + from benchflow.sandbox.daytona import DaytonaSandbox + + monkeypatch.setenv("BENCHFLOW_DAYTONA_ACP_TRANSPORT", "ssh") + env = _stub_sandbox() + with ( + patch.object( + process_pkg.DaytonaPtyProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as pty, + patch.object( + process_pkg.DaytonaProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as ssh, + ): + await DaytonaSandbox.live_process(env, agent="openhands") + + ssh.assert_awaited_once_with(env) + pty.assert_not_awaited() + + @pytest.mark.asyncio + async def test_invalid_transport_falls_back_to_pty(self, monkeypatch) -> None: + """Guards PR #921 against invalid transport config disabling Daytona.""" + from benchflow.sandbox.daytona import DaytonaSandbox + + monkeypatch.setenv("BENCHFLOW_DAYTONA_ACP_TRANSPORT", "invalid") + env = _stub_sandbox() + with ( + patch.object( + process_pkg.DaytonaPtyProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as pty, + patch.object( + process_pkg.DaytonaProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as ssh, + ): + await DaytonaSandbox.live_process(env, agent="openhands") + + pty.assert_awaited_once_with(env) + ssh.assert_not_awaited() + + @pytest.mark.asyncio + async def test_gemini_uses_ssh_transport(self) -> None: + """Guards the Gemini regression introduced by PR #896's PTY migration.""" + from benchflow.sandbox.daytona import DaytonaSandbox + + env = _stub_sandbox() + with ( + patch.object( + process_pkg.DaytonaPtyProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as pty, + patch.object( + process_pkg.DaytonaProcess, + "from_sandbox_env", + new_callable=AsyncMock, + ) as ssh, + ): + await DaytonaSandbox.live_process(env, agent="gemini") + + ssh.assert_awaited_once_with(env) + pty.assert_not_awaited() + + +class TestSingleTransportBackends: + @pytest.mark.asyncio + async def test_apple_container_uses_native_transport(self) -> None: + """Guards PR #936 against treating Apple Container as Daytona.""" + from benchflow.sandbox.apple_container import AppleContainerSandbox + + env = _stub_sandbox(_container_name="bf_run") + result = await AppleContainerSandbox.live_process(env) + + assert isinstance(result, process_pkg.AppleContainerProcess) + + @pytest.mark.asyncio + async def test_docker_uses_compose_exec_transport(self) -> None: + """Docker runs the agent through `docker compose exec -i`.""" + from benchflow.sandbox.docker import DockerSandbox + + env = _stub_sandbox() + with patch.object( + process_pkg.DockerProcess, "from_sandbox_env", return_value="docker-proc" + ) as docker: + result = await DockerSandbox.live_process(env) + + docker.assert_called_once_with(env) + assert result == "docker-proc" + + @pytest.mark.asyncio + async def test_agentcore_uses_shell_websocket_transport(self) -> None: + """AgentCore hosts the agent on its runtime-session shell WebSocket.""" + from benchflow.sandbox.agentcore import AgentCoreSandbox + + env = _stub_sandbox( + runtime_arn="arn:aws:bedrock-agentcore:us-west-2:1:runtime/x", + runtime_session_id="s" * 40, + region="us-west-2", + ) + result = await AgentCoreSandbox.live_process(env) + + assert isinstance(result, process_pkg.AgentCoreProcess) + + @pytest.mark.asyncio + async def test_backend_without_transport_fails_actionably(self) -> None: + """A backend with no live transport must say so, not borrow another's. + + Modal previously fell through the ACP layer's ``else`` branch and was + handed a ``DaytonaProcess``, which failed deep inside Daytona SSH setup + with an unrelated error. + """ + + class _NoTransportSandbox(BaseSandbox): + def __init__(self) -> None: # no BaseSandbox init needed here + pass + + def _validate_definition(self) -> None: ... + + @classmethod + def preflight(cls) -> None: ... + + async def start(self, force_build: bool) -> None: ... + async def stop(self, delete: bool) -> None: ... + async def upload_file(self, source_path, target_path) -> None: ... + async def upload_dir( + self, source_dir, target_dir, service="main" + ) -> None: ... + async def download_file(self, source_path, target_path) -> None: ... + async def download_dir( + self, source_dir, target_dir, service="main" + ) -> None: ... + async def exec(self, command, **kwargs): ... + + with pytest.raises(NotImplementedError) as excinfo: + await _NoTransportSandbox().live_process() + + assert "does not provide a live agent transport" in str(excinfo.value) diff --git a/tests/test_sandbox_provider_registry_drift.py b/tests/test_sandbox_provider_registry_drift.py index add1ca0ac..aed9aff51 100644 --- a/tests/test_sandbox_provider_registry_drift.py +++ b/tests/test_sandbox_provider_registry_drift.py @@ -31,17 +31,22 @@ def test_registry_is_the_single_source_of_truth() -> None: # Locks the current set + docker-first order; adding a provider is then a # deliberate edit here + a test update, never a silent scatter. - assert SANDBOX_PROVIDERS == ("docker", "daytona", "modal", "apple-container") + assert SANDBOX_PROVIDERS == ( + "docker", + "daytona", + "modal", + "apple-container", + "agentcore", + ) assert frozenset(SANDBOX_PROVIDERS) == SANDBOX_PROVIDER_SET def test_providers_phrase_is_byte_identical() -> None: # The refactor must be behavior-preserving for every help/error string that # used to hand-write this phrase. - assert providers_phrase() == "docker, daytona, modal, or apple-container" - assert ( - providers_phrase(quote=True) - == "'docker', 'daytona', 'modal', or 'apple-container'" + assert providers_phrase() == "docker, daytona, modal, apple-container, or agentcore" + assert providers_phrase(quote=True) == ( + "'docker', 'daytona', 'modal', 'apple-container', or 'agentcore'" ) @@ -49,10 +54,10 @@ def test_model_proxy_placement_is_explicit_for_every_provider() -> None: """Guards PR #936 against routing host loopback into an Apple VM.""" assert PROVIDERS_BY_NAME["docker"].model_proxy is ModelProxyLocation.HOST - for provider in ("daytona", "modal", "apple-container"): + for provider in ("daytona", "modal", "apple-container", "agentcore"): assert PROVIDERS_BY_NAME[provider].model_proxy is ModelProxyLocation.SANDBOX assert ( - frozenset({"daytona", "modal", "apple-container"}) + frozenset({"daytona", "modal", "apple-container", "agentcore"}) == SANDBOX_MODEL_PROXY_PROVIDERS ) assert OFF_BOX_MODEL_PROVIDERS is SANDBOX_MODEL_PROXY_PROVIDERS @@ -98,6 +103,7 @@ def test_optional_extras_match_pyproject() -> None: assert OPTIONAL_SANDBOX_EXTRAS == { "daytona": "sandbox-daytona", "modal": "sandbox-modal", + "agentcore": "sandbox-agentcore", } diff --git a/uv.lock b/uv.lock index dec20d7f7..98ae8e97f 100644 --- a/uv.lock +++ b/uv.lock @@ -308,6 +308,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/df/73/b6e24bd22e6720ca8ee9a85a0c4a2971af8497d8f3193fa05390cbd46e09/backoff-2.2.1-py3-none-any.whl", hash = "sha256:63579f9a0628e06278f7e47b7d7d5b6ce20dc65c5e96a6f3ca99a6adca0396e8", size = 15148, upload-time = "2022-10-05T19:19:30.546Z" }, ] +[[package]] +name = "bedrock-agentcore" +version = "1.18.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "boto3" }, + { name = "botocore" }, + { name = "pydantic" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "urllib3" }, + { name = "uvicorn" }, + { name = "websockets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/e8/8149bb168ef3825884685eb98ef868f5bb5f796e4ea58eb98e8df3248937/bedrock_agentcore-1.18.1.tar.gz", hash = "sha256:ed4aa8822e39aa5846b9d68d1461afee19f2f795c1f6d71a61bc650ecc00ee31", size = 1030729, upload-time = "2026-07-17T19:55:21.354Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/fe/26cc0a66920035d2740c062bd46df252b15708b7723e13ccae0ab7b334db/bedrock_agentcore-1.18.1-py3-none-any.whl", hash = "sha256:6d8001f09cdfca0a20650ea691426149f521037d8d56917057fc191c00801086", size = 450684, upload-time = "2026-07-17T19:55:19.334Z" }, +] + [[package]] name = "benchflow" version = "0.6.5" @@ -345,6 +364,10 @@ judge = [ { name = "google-genai" }, { name = "openai" }, ] +sandbox-agentcore = [ + { name = "bedrock-agentcore" }, + { name = "boto3" }, +] sandbox-daytona = [ { name = "daytona" }, { name = "tenacity" }, @@ -366,7 +389,9 @@ requires-dist = [ { name = "agent-client-protocol", specifier = ">=0.10" }, { name = "anthropic", marker = "extra == 'judge'", specifier = ">=0.40" }, { name = "anyio", specifier = ">=4.0" }, + { name = "bedrock-agentcore", marker = "extra == 'sandbox-agentcore'", specifier = ">=1.18" }, { name = "boto3", marker = "extra == 'bedrock'", specifier = ">=1.40" }, + { name = "boto3", marker = "extra == 'sandbox-agentcore'", specifier = ">=1.43.31" }, { name = "datasets", marker = "extra == 'trl'", specifier = ">=2.19.0" }, { name = "daytona", marker = "extra == 'sandbox-daytona'", specifier = ">=0.184.0" }, { name = "deepagents", marker = "extra == 'deepagents'", specifier = ">=0.6" }, @@ -392,34 +417,34 @@ requires-dist = [ { name = "ty", marker = "extra == 'dev'", specifier = ">=0.0.1a1" }, { name = "typer", specifier = ">=0.9" }, ] -provides-extras = ["dev", "train", "trl", "sandbox-daytona", "sandbox-modal", "bedrock", "judge", "deepagents"] +provides-extras = ["dev", "train", "trl", "sandbox-daytona", "sandbox-modal", "sandbox-agentcore", "bedrock", "judge", "deepagents"] [[package]] name = "boto3" -version = "1.43.6" +version = "1.43.56" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore" }, { name = "jmespath" }, { name = "s3transfer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0a/37/78c630d1308964aa9abf44951d9c4df776546ff37251ec2434944e205c4e/boto3-1.43.6.tar.gz", hash = "sha256:e6315effaf12b890b99956e6f8e2c3000a3f64e4ee91943cec3895ce9a836afb", size = 113153, upload-time = "2026-05-07T20:49:59.694Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/05/23e1aa8c9e4b0399a61e7fd65c4f9cc0625121f24760e37471f776404abb/boto3-1.43.56.tar.gz", hash = "sha256:57c90df9fb026f2e6ae22530861198130203733c5c9ec4e5cca3a4037f5a8db4", size = 112673, upload-time = "2026-07-24T19:31:48.606Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/e2/3c2eef44f55eafab256836d1d9479bd6a74f70c26cbfdc0639a0e23e4327/boto3-1.43.6-py3-none-any.whl", hash = "sha256:179601ec2992726a718053bf41e43c223ceba397d31ceab11f64d9c910d9fc3a", size = 140502, upload-time = "2026-05-07T20:49:57.8Z" }, + { url = "https://files.pythonhosted.org/packages/b8/57/3a960c9f581c00f2a591901b46e035ff79ab3956d16607f12306b3b8d483/boto3-1.43.56-py3-none-any.whl", hash = "sha256:feb699d4ab241ef5c1b80bb58277be2aaad365cd4b672d7817e0bc59ee45131b", size = 140026, upload-time = "2026-07-24T19:31:47.155Z" }, ] [[package]] name = "botocore" -version = "1.43.6" +version = "1.43.56" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jmespath" }, { name = "python-dateutil" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/79/a7/23d0f5028011455096a1eeac0ddf3cbe147b3e855e127342f8202552194d/botocore-1.43.6.tar.gz", hash = "sha256:b1e395b347356860398da42e61c808cf1e34b6fa7180cf2b9d87d986e1a06ba0", size = 15336070, upload-time = "2026-05-07T20:49:48.14Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/cc/7f84a5d3071fe878380e9f610ab36ca87b8cbbc4aa81ba2727f90e1f3ea3/botocore-1.43.56.tar.gz", hash = "sha256:6c01f85f0ff9863076f4c761e74ee3aa96c5ccc1ad09fc1efd62ef8f2d22bf57", size = 15733117, upload-time = "2026-07-24T19:31:38.125Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e5/c8/6f47223840e8d8cfa8c9f7c0ec1b77970417f257fc885169ff4f6326ce09/botocore-1.43.6-py3-none-any.whl", hash = "sha256:b6d1fdbc6f65a5fe0b7e947823aa37535d3f39f3ba4d21110fab1f55bbbcc04b", size = 15017094, upload-time = "2026-05-07T20:49:44.964Z" }, + { url = "https://files.pythonhosted.org/packages/5c/cd/86fe9e659e9699f62f8dd5ecd8c6725474334b23cab8aa71d82b5f56f1a4/botocore-1.43.56-py3-none-any.whl", hash = "sha256:aafc741f1b10f6fd63253eaf6ea029680c1ff436d87e1b8969d62aefa0c76976", size = 15418773, upload-time = "2026-07-24T19:31:34.758Z" }, ] [[package]] @@ -3667,14 +3692,14 @@ wheels = [ [[package]] name = "s3transfer" -version = "0.17.0" +version = "0.19.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9b/ec/7c692cde9125b77e84b307354d4fb705f98b8ccad59a036d5957ca75bfc3/s3transfer-0.17.0.tar.gz", hash = "sha256:9edeb6d1c3c2f89d6050348548834ad8289610d886e5bf7b7207728bd43ce33a", size = 155337, upload-time = "2026-04-29T22:07:36.33Z" } +sdist = { url = "https://files.pythonhosted.org/packages/76/43/35e4d8aa320bffe8287fe8f65f578fa2d2db0a64212f0e710dce58267854/s3transfer-0.19.2.tar.gz", hash = "sha256:ba0309fd86be3c27dbf78cdd813c13c5e1df16e5874b99d2535ebbdfb9892993", size = 165592, upload-time = "2026-07-22T19:30:44.432Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/87/72/c6c32d2b657fa3dad1de340254e14390b1e334ce38268b7ad51abda3c8c2/s3transfer-0.17.0-py3-none-any.whl", hash = "sha256:ce3801712acf4ad3e89fb9990df97b4972e93f4b3b0004d214be5bce12814c20", size = 86811, upload-time = "2026-04-29T22:07:34.966Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl", hash = "sha256:d8168eccca828cbb2cd573675333f3bddd254313a9c42494b84c76b539e8ba25", size = 90216, upload-time = "2026-07-22T19:30:43.251Z" }, ] [[package]] From ee004a16cb998a67bf6a8039859ddf343ee7f41d Mon Sep 17 00:00:00 2001 From: Bingran You Date: Mon, 27 Jul 2026 13:39:15 -0700 Subject: [PATCH 02/13] fix(agentcore): share runtimes across rollouts so a matrix can fan out MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Makes AgentCore usable as a parallel backend and fixes a bug that would silently corrupt any multi-trial run. A rollout is a **session**, not a runtime. One registered runtime hosts many concurrent, filesystem-isolated microVMs (verified: 8 concurrent sessions on one runtime), and the account quotas are lopsided in exactly that direction — 5000 Active Session Workloads against 100 Total Agents, with CreateAgentRuntime and ListAgentRuntimes both limited to 5/s. The first cut got this backwards: it created and deleted a runtime per rollout and named it after the *task*. Three trials of one task therefore raced to create the same runtime, and whichever finished first deleted it out from under the other two. Two different tasks ran fine, which is why the initial testing missed it. Runtime identity now follows the image's content digest, computed from the build context before anything is built, and provisioning is single-flighted per key in a new `agentcore_provisioning` module. N concurrent rollouts of one task now perform exactly one build, one push, and one registration, then open N sessions. Verified live: 5 concurrent rollouts of one task produced 1 runtime and 5 isolated sessions in 17.7s, and the shared runtime survived teardown. Failures are deliberately not memoized so a transient throttle cannot poison the rest of a long run. The digest hashes file *contents*, not paths or mtimes: BenchFlow copies tasks into temp directories, so any location-based identity would change every run and defeat image reuse entirely. Because runtimes are now shared they outlive the run, like a built Docker image, and `bench sandbox cleanup` reclaims them by age. Only runtimes tagged benchflow-managed are touched, and unreadable tags fail closed — verified live against a runtime created outside BenchFlow, which was correctly skipped. Also fixed while wiring cleanup up: - `bench sandbox cleanup` aborted entirely when the Daytona SDK was installed but had no credentials, stranding AgentCore runtimes against a 100-per- account quota. Each backend's cleanup is now isolated. - Cleanup is gated on BENCHFLOW_AGENTCORE_ROLE_ARN rather than on an importable boto3, which several unrelated extras pull in. Without that gate the command reached out to AWS on any machine with credentials configured — including from the test suite. Adds an image-size preflight for AgentCore's hard, non-adjustable 2 GB limit, so an oversized environment fails with a message naming the quota instead of an opaque runtime error later. --- docs/running-benchmarks.md | 34 ++ src/benchflow/cli/sandbox.py | 84 ++++- src/benchflow/sandbox/agentcore.py | 267 +++++++------ .../sandbox/agentcore_provisioning.py | 159 ++++++++ src/benchflow/sandbox/agentcore_reaper.py | 103 +++++ tests/test_agentcore_parallel.py | 356 ++++++++++++++++++ 6 files changed, 876 insertions(+), 127 deletions(-) create mode 100644 src/benchflow/sandbox/agentcore_provisioning.py create mode 100644 src/benchflow/sandbox/agentcore_reaper.py create mode 100644 tests/test_agentcore_parallel.py diff --git a/docs/running-benchmarks.md b/docs/running-benchmarks.md index 1988506f3..4546b0a23 100644 --- a/docs/running-benchmarks.md +++ b/docs/running-benchmarks.md @@ -393,6 +393,40 @@ Sessions default to a 15-minute idle timeout and an 8-hour lifetime; override with `BENCHFLOW_AGENTCORE_IDLE_TIMEOUT_SEC` / `BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC` if agent turns are long enough to risk reclamation mid-run. +**Running a large matrix.** A rollout is a *session*, not a runtime: one +registered runtime hosts many concurrent, filesystem-isolated microVMs. The +account quotas point the same way — 5,000 concurrent *Active Session Workloads* +against only 100 *Total Agents*, with `CreateAgentRuntime` limited to 5/s. So +BenchFlow builds and registers once per distinct task image (keyed by a digest +of the build context) and every trial and skill arm of that image opens another +session against the shared runtime. Concurrency is therefore bounded by +sessions, not by how many images you have. + +Two quotas to check before a big run: + +| Quota | Default | Adjustable | +|---|---|---| +| Active Session Workloads per account | 5,000 | yes | +| Total Agents (runtimes) per account | 100 | yes | +| **Max image size** | **2 GB** | **no** | + +A full 88-task × 2-skill-arm matrix is 176 distinct images, over the default +100-runtime cap — request an increase first. The 2 GB image limit is a hard +service quota: heavy environments (LaTeX, Playwright, large model snapshots) +cannot run here at all, and BenchFlow fails the build with a message naming the +cap rather than letting it surface later as an opaque runtime error. Run those +tasks on Docker or Daytona. + +Runtimes are shared, so they deliberately outlive the run that created them — +like a built Docker image. Reclaim them by age: + +```bash +bench sandbox cleanup --max-age 1440 +``` + +Only runtimes tagged `benchflow-managed` are ever deleted, and cleanup is a +no-op unless `BENCHFLOW_AGENTCORE_ROLE_ARN` is set. + For large-scale runs (100+ tasks), use Daytona or Modal with high concurrency: ```bash diff --git a/src/benchflow/cli/sandbox.py b/src/benchflow/cli/sandbox.py index 20fb3dd44..53496412c 100644 --- a/src/benchflow/cli/sandbox.py +++ b/src/benchflow/cli/sandbox.py @@ -124,24 +124,82 @@ def sandbox_list_local() -> None: console.print(f"\n[bold]{total} sandbox(es)[/bold]") +def _cleanup_agentcore_runtimes(*, dry_run: bool, max_age_minutes: int) -> bool: + """Reap stale AgentCore runtimes. False when AgentCore is not in play. + + AgentCore runtimes are shared across rollouts and intentionally outlive the + run that created them, so unlike Docker they accumulate — and *Total Agents + per Account* defaults to only 100. + + Gated on ``BENCHFLOW_AGENTCORE_ROLE_ARN`` rather than merely on an + importable SDK. ``boto3`` arrives with several unrelated extras, so without + this gate ``bench sandbox cleanup`` would reach out to AWS on any machine + that happens to have credentials configured — including during tests. + """ + import os + + if not os.environ.get("BENCHFLOW_AGENTCORE_ROLE_ARN"): + return False + try: + import boto3 + + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + except ImportError: + return False + + region = os.environ.get("BENCHFLOW_AGENTCORE_REGION") or "us-west-2" + try: + control = boto3.Session(region_name=region).client("bedrock-agentcore-control") + report = reap_stale_runtimes( + control, max_age_minutes=max_age_minutes, dry_run=dry_run + ) + except Exception as exc: + # Missing/expired credentials are a "nothing to do here", not a crash + # in the middle of a cleanup that may still have Daytona work to do. + print_error(f"Skipping AgentCore cleanup: {exc}") + return False + + verb = "Would delete" if dry_run else "Deleted" + console.print(f"AgentCore ({region}): {report.summary()}") + if report.deleted: + console.print(f" {verb}: {', '.join(report.deleted)}") + return True + + def sandbox_cleanup(*, dry_run: bool, max_age_minutes: int) -> None: - """Clean up orphaned Daytona sandboxes. + """Clean up orphaned Daytona sandboxes and stale AgentCore runtimes. - Like ``sandbox list``, this is a no-op (not an error) when the optional - Daytona SDK is absent: only Daytona has persistent sandboxes to reap. + Both are opt-in extras, so a missing SDK is a no-op rather than an error; + only these two backends leave anything behind between runs. Docker + sandboxes are torn down per run. """ - if not _daytona_sdk_available(): + cleaned_any = False + + if _daytona_sdk_available(): + from benchflow.cli import main as cli_main + + try: + cli_main._cleanup_daytona_sandboxes( + dry_run=dry_run, max_age_minutes=max_age_minutes + ) + cleaned_any = True + except (Exception, typer.Exit) as exc: + # The Daytona SDK is installed but unusable (typically no + # DAYTONA_API_KEY). Report and continue: a second backend may still + # have resources to reclaim, and aborting here would silently leave + # them behind — AgentCore runtimes consume a 100-per-account quota. + print_error(f"Skipping Daytona cleanup: {exc}") + + if _cleanup_agentcore_runtimes(dry_run=dry_run, max_age_minutes=max_age_minutes): + cleaned_any = True + + if not cleaned_any: console.print( - "Nothing to clean up. The Daytona SDK is not installed " - "([cyan]uv sync --extra sandbox-daytona[/cyan]); only Daytona has " - "persistent sandboxes to reap. Docker sandboxes are torn down per run." + "Nothing to clean up. Neither the Daytona SDK " + "([cyan]uv sync --extra sandbox-daytona[/cyan]) nor the AgentCore " + "extra ([cyan]uv sync --extra sandbox-agentcore[/cyan]) is " + "installed; only those backends leave resources between runs." ) - return - from benchflow.cli import main as cli_main - - cli_main._cleanup_daytona_sandboxes( - dry_run=dry_run, max_age_minutes=max_age_minutes - ) def register_sandbox(app: typer.Typer) -> None: diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py index 77bb6cf54..e8bb3e068 100644 --- a/src/benchflow/sandbox/agentcore.py +++ b/src/benchflow/sandbox/agentcore.py @@ -5,11 +5,18 @@ pushed to ECR and registered as an *agent runtime* (a control-plane resource with its own ARN), after which sessions are invoked against that ARN. -BenchFlow therefore provisions per task: build the task's image for -``linux/arm64``, push it to ECR under a tag derived from the image's content -digest, and reuse (or create) the agent runtime bound to that tag. The digest -tag is what makes repeat runs cheap — the second task with an unchanged -environment reuses both the ECR layer set and the runtime. +**A rollout is a session, not a runtime.** One registered runtime hosts many +concurrent sessions, each an isolated microVM with its own filesystem, and the +account quotas are lopsided in exactly that direction: 5000 concurrent +sessions against 100 total runtimes, with ``CreateAgentRuntime`` limited to +5/s. So the image and the runtime are built **once per distinct task image** +— keyed by a digest of the build context, and shared by every trial and skill +arm that resolves to it — while sessions are what scale out. That sharing is +what makes this backend usable for a large parallel matrix; see +``agentcore_provisioning`` for the single-flight machinery. + +Runtimes therefore outlive the rollout that first needed them, like a built +Docker image, and are reclaimed by age via ``bench sandbox cleanup``. Two facts about the platform shape this module, both established by direct experiment against the live service rather than from the docs: @@ -38,10 +45,8 @@ import asyncio import base64 -import hashlib import io import os -import re import shlex import shutil import subprocess @@ -52,6 +57,7 @@ from typing import TYPE_CHECKING, Any from benchflow._paths import iter_safe_tree +from benchflow.sandbox import agentcore_provisioning as provisioning from benchflow.sandbox._base import BaseSandbox, ExecResult, wrap_command_with_env_file from benchflow.sandbox.protocol import SandboxStartupError from benchflow.task.config import SandboxConfig @@ -72,7 +78,6 @@ # and the wrapper adds shell scaffolding, so keep a wide margin. _MAX_INLINE_UPLOAD_BYTES = 24 * 1024 _MAX_DOWNLOAD_BYTES = 64 * 1024 * 1024 -_ECR_TAG_RE = re.compile(r"[^a-zA-Z0-9_.-]+") # ``_PING_SHIM`` implements the plain HTTP contract (GET /ping, # POST /invocations), so register the runtime as HTTP rather than relying on # whatever the service currently defaults to. @@ -162,7 +167,6 @@ def __init__( self.runtime_arn: str | None = None self.runtime_session_id: str | None = None self._runtime_id: str | None = None - self._created_runtime = False self._account_id: str | None = None self._clients: dict[str, Any] = {} self._client_lock = threading.Lock() @@ -248,10 +252,6 @@ def _resolve_account_id(self) -> str: def _ecr_registry(self) -> str: return f"{self._resolve_account_id()}.dkr.ecr.{self.region}.amazonaws.com" - def _image_tag(self) -> str: - safe = _ECR_TAG_RE.sub("_", self.environment_name)[:80] - return f"bf-{safe}".lower() - async def start(self, force_build: bool) -> None: try: image_uri = await self._publish_image(force_build=force_build) @@ -275,13 +275,41 @@ async def start(self, force_build: bool) -> None: build_timeout_sec=self.task_env_config.build_timeout_sec, ) from exc + def _image_identity(self) -> tuple[str, str]: + """``(context digest, ECR tag)`` for this task's image. + + Computed from the build context's *contents* before anything is built, + so concurrent rollouts of the same task agree on the tag without each + running a build to discover it. + """ + digest = provisioning.build_context_digest( + self.environment_dir, self._generated_dockerfile_text() + ) + return digest, provisioning.image_tag(self.environment_name, digest) + async def _publish_image(self, *, force_build: bool) -> str: - """Build the task image with the ping shim and push it to ECR.""" - build_context, dockerfile, shim = self._materialize_build_context() - registry = self._ecr_registry() - repo = self._ecr_repository - local_tag = self._image_tag() + """Build the task image with the ping shim and push it to ECR. + Deduplicated process-wide by content digest: a fan-out of rollouts over + one task builds and pushes exactly once, and every later run of an + unchanged task skips the build entirely on the ECR existence check. + """ + _digest, tag = self._image_identity() + image_uri = f"{self._ecr_registry()}/{self._ecr_repository}:{tag}" + cache_key = f"image:{image_uri}:{force_build}" + + async def _publish() -> str: + await asyncio.to_thread(self._ensure_ecr_repository) + if not force_build and await asyncio.to_thread(self._ecr_image_exists, tag): + self.logger.info("Reusing published AgentCore image %s", image_uri) + return image_uri + await self._build_and_push(image_uri, force_build=force_build) + return image_uri + + return await provisioning.once(cache_key, _publish) + + async def _build_and_push(self, image_uri: str, *, force_build: bool) -> None: + build_context, dockerfile, shim = self._materialize_build_context() build_args = [ "docker", "build", @@ -290,7 +318,7 @@ async def _publish_image(self, *, force_build: bool) -> str: "-f", str(dockerfile), "-t", - local_tag, + image_uri, ] if force_build: build_args.append("--no-cache") @@ -312,28 +340,10 @@ async def _publish_image(self, *, force_build: bool) -> str: f"{(result.stderr or result.stdout or '').strip()[:4000]}" ) - # Tag by image content digest so an unchanged environment reuses the - # already-pushed layers and the already-registered agent runtime. - digest = await asyncio.to_thread( - _run, "docker", "images", "--no-trunc", "--quiet", local_tag - ) - content_id = (digest.stdout or "").strip().replace("sha256:", "")[:32] - if not content_id: - content_id = hashlib.sha256(local_tag.encode()).hexdigest()[:32] - remote_tag = f"{local_tag}-{content_id}" - image_uri = f"{registry}/{repo}:{remote_tag}" - - await asyncio.to_thread(self._ensure_ecr_repository) - if not force_build and await asyncio.to_thread( - self._ecr_image_exists, remote_tag - ): - self.logger.info("Reusing published AgentCore image %s", image_uri) - return image_uri + await self._reject_oversized_image(image_uri) + registry = self._ecr_registry() await asyncio.to_thread(self._docker_login, registry) - retag = await asyncio.to_thread(_run, "docker", "tag", local_tag, image_uri) - if retag.returncode != 0: - raise RuntimeError(f"docker tag failed: {(retag.stderr or '')[:2000]}") push = await asyncio.to_thread(_run, "docker", "push", image_uri, timeout=1800) if push.returncode != 0: raise RuntimeError( @@ -341,7 +351,21 @@ async def _publish_image(self, *, force_build: bool) -> str: f"{(push.stderr or push.stdout or '').strip()[:4000]}" ) self.logger.info("Pushed AgentCore image %s", image_uri) - return image_uri + + async def _reject_oversized_image(self, image_uri: str) -> None: + """Fail before push when the image exceeds AgentCore's hard 2 GB cap.""" + inspect = await asyncio.to_thread( + _run, "docker", "image", "inspect", "-f", "{{.Size}}", image_uri + ) + raw = (inspect.stdout or "").strip() + if inspect.returncode != 0 or not raw.isdigit(): + # Can't measure — let the push proceed rather than block on a + # diagnostic that failed for an unrelated reason. + self.logger.debug("Could not measure image size for %s", image_uri) + return + message = provisioning.image_size_error(int(raw), image_uri) + if message: + raise SandboxStartupError(message, sandbox_id=image_uri) def _materialize_build_context(self) -> tuple[Path, Path, Path]: """Return (context_dir, dockerfile, shim) for an image with the shim. @@ -352,27 +376,34 @@ def _materialize_build_context(self) -> tuple[Path, Path, Path]: transient — the caller removes them once the build finishes. """ context = self.environment_dir - base_dockerfile = context / "Dockerfile" + shim_path = context / provisioning.GENERATED_SHIM + shim_path.write_text(_PING_SHIM) + generated = context / provisioning.GENERATED_DOCKERFILE + generated.write_text(self._generated_dockerfile_text()) + return context, generated, shim_path + + def _generated_dockerfile_text(self) -> str: + """The Dockerfile BenchFlow actually builds, as text. + + Kept separate from writing it so the image digest can be computed from + the same bytes without touching the task tree. + """ if self.task_env_config.docker_image: base = f"FROM {self.task_env_config.docker_image}\n" else: - base = base_dockerfile.read_text() - - shim_path = context / ".benchflow_agentcore_shim.py" - shim_path.write_text(_PING_SHIM) - generated = context / "Dockerfile.benchflow-agentcore" - generated.write_text( + base = (self.environment_dir / "Dockerfile").read_text() + return ( base + "\n" + "# --- BenchFlow AgentCore runtime contract ---\n" + "# AgentCore refuses command execution for a session whose\n" + "# container does not answer GET /ping on :8080.\n" - + "COPY .benchflow_agentcore_shim.py /opt/benchflow_agentcore_shim.py\n" + + f"COPY {provisioning.GENERATED_SHIM} " + + "/opt/benchflow_agentcore_shim.py\n" + "EXPOSE 8080\n" + "ENTRYPOINT []\n" + 'CMD ["python3", "/opt/benchflow_agentcore_shim.py"]\n' ) - return context, generated, shim_path def _ensure_ecr_repository(self) -> None: from botocore.exceptions import ClientError @@ -409,54 +440,77 @@ def _docker_login(self, registry: str) -> None: if proc.returncode != 0: raise RuntimeError(f"ECR docker login failed: {(proc.stderr or '')[:1000]}") - def _runtime_name(self) -> str: - # AgentCore runtime names allow [A-Za-z][A-Za-z0-9_]*. - safe = re.sub(r"[^a-zA-Z0-9_]+", "_", f"bf_{self.environment_name}") - return safe[:48] - async def _ensure_runtime(self, image_uri: str) -> str: - """Reuse the agent runtime bound to *image_uri*, or create it.""" + """Resolve the shared agent runtime for *image_uri*, creating it once. + + The runtime is named after the image's content digest, so every rollout + of that image — each trial, and both skill arms when their images match + — resolves to the same runtime and simply opens another session against + it. Keying on the task name instead meant concurrent trials raced to + create one runtime and the first to finish deleted it while the others + were still running. + """ + digest, _tag = self._image_identity() + name = provisioning.runtime_name(self.environment_name, digest) + + async def _create() -> tuple[str, str]: + return await self._create_or_adopt_runtime(name, image_uri) + + arn, runtime_id = await provisioning.once(f"runtime:{name}", _create) + self._runtime_id = runtime_id + return arn + + async def _create_or_adopt_runtime( + self, name: str, image_uri: str + ) -> tuple[str, str]: + """Create the runtime, or adopt an equivalent one that already exists. + + Attempting the create first and falling back on conflict keeps the + common path off ``ListAgentRuntimes``, whose 5/s quota would otherwise + throttle a large matrix run before it started. + """ + from botocore.exceptions import ClientError + control = self._client("bedrock-agentcore-control") - name = self._runtime_name() - - existing = await asyncio.to_thread(self._find_runtime, control, name) - if existing is not None: - arn, runtime_id, current_image = existing - self._runtime_id = runtime_id - if current_image == image_uri: - await asyncio.to_thread(self._wait_ready, control, runtime_id) - self.logger.info("Reusing AgentCore runtime %s", arn) - return arn - await asyncio.to_thread( - control.update_agent_runtime, - agentRuntimeId=runtime_id, + try: + created = await asyncio.to_thread( + control.create_agent_runtime, + agentRuntimeName=name, agentRuntimeArtifact={ "containerConfiguration": {"containerUri": image_uri} }, roleArn=self._require_role_arn(), networkConfiguration={"networkMode": "PUBLIC"}, protocolConfiguration=_PROTOCOL_CONFIGURATION, + lifecycleConfiguration=self._lifecycle_configuration(), + description=f"BenchFlow sandbox for {self.environment_name}", + tags={ + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + "benchflow-task": self.environment_name[:120], + }, + ) + except ClientError as exc: + code = exc.response["Error"]["Code"] + if code not in {"ConflictException", "ValidationException"}: + raise + existing = await asyncio.to_thread( + provisioning.find_runtime_by_name, control, name ) + if existing is None: + raise + arn, runtime_id, _image = existing await asyncio.to_thread(self._wait_ready, control, runtime_id) - self.logger.info("Updated AgentCore runtime %s to %s", arn, image_uri) - return arn - - created = await asyncio.to_thread( - control.create_agent_runtime, - agentRuntimeName=name, - agentRuntimeArtifact={ - "containerConfiguration": {"containerUri": image_uri} - }, - roleArn=self._require_role_arn(), - networkConfiguration={"networkMode": "PUBLIC"}, - protocolConfiguration=_PROTOCOL_CONFIGURATION, - lifecycleConfiguration=self._lifecycle_configuration(), - description=f"BenchFlow sandbox for {self.environment_name}", + self.logger.info("Adopted existing AgentCore runtime %s", arn) + return arn, runtime_id + + runtime_id = created["agentRuntimeId"] + await asyncio.to_thread(self._wait_ready, control, runtime_id) + self.logger.info( + "Registered AgentCore runtime %s for %s", + created["agentRuntimeArn"], + image_uri, ) - self._created_runtime = True - self._runtime_id = created["agentRuntimeId"] - await asyncio.to_thread(self._wait_ready, control, created["agentRuntimeId"]) - return created["agentRuntimeArn"] + return created["agentRuntimeArn"], runtime_id def _lifecycle_configuration(self) -> dict[str, int]: """Session idle/lifetime caps. @@ -484,19 +538,6 @@ def _require_role_arn(self) -> str: "can assume and that can pull from ECR and write CloudWatch logs." ) - @staticmethod - def _find_runtime(control: Any, name: str) -> tuple[str, str, str | None] | None: - paginator_pages = control.list_agent_runtimes() - for runtime in paginator_pages.get("agentRuntimes", []): - if runtime.get("agentRuntimeName") != name: - continue - runtime_id = runtime["agentRuntimeId"] - detail = control.get_agent_runtime(agentRuntimeId=runtime_id) - artifact = detail.get("agentRuntimeArtifact") or {} - image = (artifact.get("containerConfiguration") or {}).get("containerUri") - return detail["agentRuntimeArn"], runtime_id, image - return None - @staticmethod def _wait_ready(control: Any, runtime_id: str) -> None: import time @@ -573,9 +614,20 @@ async def _warm_session(self) -> None: ) async def stop(self, delete: bool) -> None: - await self._safe_teardown(delete_runtime=delete) + """End this rollout's session. The runtime is shared and outlives it. + + A rollout owns a *session*, not a runtime: concurrent trials of the + same task run as separate sessions against one registered runtime, so + deleting the runtime here would tear down sandboxes still in use — and + would also burn the 5/s ``DeleteAgentRuntime``/``CreateAgentRuntime`` + quotas re-registering it for the next rollout. - async def _safe_teardown(self, *, delete_runtime: bool = True) -> None: + Runtimes are therefore left in place, like a built Docker image, and + reclaimed by age with ``bench sandbox cleanup --sandbox agentcore``. + """ + await self._safe_teardown() + + async def _safe_teardown(self) -> None: if self.runtime_arn and self.runtime_session_id: try: await asyncio.to_thread( @@ -588,19 +640,6 @@ async def _safe_teardown(self, *, delete_runtime: bool = True) -> None: self.logger.warning("Failed to stop AgentCore session: %s", exc) self.runtime_session_id = None - if delete_runtime and self._created_runtime and self._runtime_id: - try: - await asyncio.to_thread( - self._client("bedrock-agentcore-control").delete_agent_runtime, - agentRuntimeId=self._runtime_id, - ) - self.logger.info("Deleted AgentCore runtime %s", self._runtime_id) - except Exception as exc: - self.logger.warning("Failed to delete AgentCore runtime: %s", exc) - self._runtime_id = None - self.runtime_arn = None - self._created_runtime = False - # ------------------------------------------------------------------ exec async def live_process(self, *, agent: str | None = None) -> LiveProcess: diff --git a/src/benchflow/sandbox/agentcore_provisioning.py b/src/benchflow/sandbox/agentcore_provisioning.py new file mode 100644 index 000000000..f8c518763 --- /dev/null +++ b/src/benchflow/sandbox/agentcore_provisioning.py @@ -0,0 +1,159 @@ +"""Process-wide, single-flight provisioning of AgentCore images and runtimes. + +A rollout on AgentCore is a **session**, not a runtime. One agent runtime can +host many concurrent sessions, each an isolated microVM with its own +filesystem — measured at 8 concurrent sessions on one runtime, and the account +quota for *Active Session Workloads* is 5000 against only 100 *Total Agents*. +So the expensive, rate-limited artifacts (an ECR image and a registered +runtime) must be created **once per distinct task image** and shared by every +rollout that uses it, while sessions are what scale out. + +Getting that wrong is not merely slow. Keying a runtime on the task name meant +that three trials of one task raced to create the same runtime and the first to +finish deleted it out from under the other two. Keying on the *content* of the +build context makes the mapping deterministic: same image ⇒ same runtime ⇒ one +build, one push, one registration, N sessions. + +The control plane is also far tighter than the data plane +(``CreateAgentRuntime`` and ``ListAgentRuntimes`` are 5/s, while +``InvokeAgentRuntimeCommand`` is 200/s), which is why nothing here may run per +rollout. Results are memoized for the life of the process and every miss is +funnelled through a per-key lock. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import re +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import Any + +logger = logging.getLogger("benchflow").getChild("agentcore") + +# Names generated into the build context; excluded from its digest so the +# digest describes the *task*, not our own scaffolding. +GENERATED_DOCKERFILE = "Dockerfile.benchflow-agentcore" +GENERATED_SHIM = ".benchflow_agentcore_shim.py" +_GENERATED_NAMES = frozenset({GENERATED_DOCKERFILE, GENERATED_SHIM}) + +# Tag convention shared with the Daytona reaper so cleanup can tell BenchFlow's +# resources apart from anything else in the account. +MANAGED_TAG = "benchflow-managed" +MANAGED_VALUE = "1" + +# Not adjustable per AWS service quotas: "Maximum size (in MB) for a Docker +# image in an AgentCore Runtime" = 2048. +MAX_IMAGE_MB = 2048 + +_GLOBAL_LOCK = asyncio.Lock() +_KEY_LOCKS: dict[str, asyncio.Lock] = {} +_RESULTS: dict[str, Any] = {} + + +async def once[T](key: str, factory: Callable[[], Awaitable[T]]) -> T: + """Run *factory* at most once per *key* for the life of the process. + + Concurrent callers with the same key block on one lock and then read the + memoized result, so a fan-out of N rollouts over the same task image + performs exactly one build, one push, and one runtime registration. + + Failures are deliberately not memoized: a transient throttle or a network + blip should not poison every later rollout in a long matrix run. + """ + if key in _RESULTS: + return _RESULTS[key] + async with _GLOBAL_LOCK: + lock = _KEY_LOCKS.setdefault(key, asyncio.Lock()) + async with lock: + if key in _RESULTS: + return _RESULTS[key] + value = await factory() + _RESULTS[key] = value + return value + + +def reset_cache() -> None: + """Drop memoized provisioning state (tests only).""" + _RESULTS.clear() + _KEY_LOCKS.clear() + + +def build_context_digest(context_dir: Path, dockerfile_text: str) -> str: + """Content digest of everything that determines the built image. + + Hashes file *contents* rather than paths and mtimes on purpose: BenchFlow + copies tasks into temporary directories before a run, so any identity based + on location or timestamp would change every run and defeat image reuse + entirely. Symlinks are skipped for the same reason they are skipped on + upload (#411) — they must not pull host files into the image identity. + """ + digest = hashlib.sha256() + digest.update(dockerfile_text.encode()) + for path in sorted(context_dir.rglob("*")): + if path.is_symlink() or not path.is_file(): + continue + relative = path.relative_to(context_dir).as_posix() + if relative in _GENERATED_NAMES: + continue + digest.update(relative.encode()) + digest.update(b"\0") + digest.update(hashlib.sha256(path.read_bytes()).digest()) + return digest.hexdigest() + + +def image_tag(task_name: str, digest: str) -> str: + """ECR tag for a task image: readable prefix plus content digest.""" + safe = re.sub(r"[^a-zA-Z0-9_.-]+", "-", task_name).strip("-.")[:40].lower() + return f"bf-{safe}-{digest[:16]}" + + +def runtime_name(task_name: str, digest: str) -> str: + """Agent-runtime name derived from image identity, not the task name. + + Two rollouts of the same task image — different trials, or the with-skill + and no-skill arms when their images happen to match — resolve to the same + name and therefore share one runtime. AgentCore accepts + ``[A-Za-z][A-Za-z0-9_]*``; the ``bf_`` prefix guarantees the leading letter. + """ + safe = re.sub(r"[^a-zA-Z0-9_]+", "_", task_name)[:28].strip("_") + return f"bf_{safe}_{digest[:12]}"[:48] + + +def image_size_error(size_bytes: int, image_uri: str) -> str | None: + """Return an error message if *size_bytes* exceeds AgentCore's hard cap. + + The 2 GB limit is **not adjustable**. Without this check an oversized image + fails later as an opaque runtime error, which reads as a task failure + rather than as an environment that AgentCore cannot host at all. + """ + size_mb = size_bytes / (1024 * 1024) + if size_mb <= MAX_IMAGE_MB: + return None + return ( + f"Image {image_uri} is {size_mb:.0f} MB, over AgentCore's " + f"{MAX_IMAGE_MB} MB per-image limit (a hard service quota, not " + "adjustable). Slim the task image, or run this task on the docker or " + "daytona sandbox instead." + ) + + +def find_runtime_by_name(control: Any, name: str) -> tuple[str, str, str | None] | None: + """Look up an existing runtime by name → ``(arn, id, image_uri)``. + + ``ListAgentRuntimes`` is a 5/s quota, so this is only ever the slow path + behind :func:`once` and the create-conflict fallback — never per rollout. + """ + paginator = control.get_paginator("list_agent_runtimes") + for page in paginator.paginate(): + for runtime in page.get("agentRuntimes", []): + if runtime.get("agentRuntimeName") != name: + continue + runtime_id = runtime["agentRuntimeId"] + detail = control.get_agent_runtime(agentRuntimeId=runtime_id) + artifact = detail.get("agentRuntimeArtifact") or {} + image = (artifact.get("containerConfiguration") or {}).get("containerUri") + return detail["agentRuntimeArn"], runtime_id, image + return None diff --git a/src/benchflow/sandbox/agentcore_reaper.py b/src/benchflow/sandbox/agentcore_reaper.py new file mode 100644 index 000000000..3abeeaf49 --- /dev/null +++ b/src/benchflow/sandbox/agentcore_reaper.py @@ -0,0 +1,103 @@ +"""Reclaim stale BenchFlow-managed AgentCore runtimes. + +Runtimes are shared across rollouts and deliberately outlive the run that +created them, so nothing deletes them on the hot path. They are also a scarce +resource: *Total Agents per Account* defaults to 100, and a full SkillsBench +matrix (one image per task per skill arm) can exceed that. Left alone, a few +large runs would exhaust the quota and every later ``CreateAgentRuntime`` +would fail. + +This mirrors ``daytona_reaper``: only resources carrying BenchFlow's managed +tag are ever considered, so a runtime someone else created in the same account +is never touched. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from typing import Any + +from benchflow.sandbox.agentcore_provisioning import MANAGED_TAG, MANAGED_VALUE + +logger = logging.getLogger("benchflow").getChild("agentcore-reaper") + +REAP_DEFAULT_MAX_AGE_MIN = 1440 + + +@dataclass +class ReapReport: + """What a reap pass considered, skipped, and deleted.""" + + scanned: int = 0 + deleted: list[str] = field(default_factory=list) + skipped_unmanaged: int = 0 + skipped_recent: int = 0 + errors: list[str] = field(default_factory=list) + + def summary(self) -> str: + return ( + f"scanned={self.scanned} deleted={len(self.deleted)} " + f"unmanaged={self.skipped_unmanaged} recent={self.skipped_recent} " + f"errors={len(self.errors)}" + ) + + +def _is_benchflow_managed(control: Any, arn: str) -> bool: + """True only when the runtime carries BenchFlow's managed tag. + + Fails closed: if the tags cannot be read, the runtime is treated as + someone else's and left alone. + """ + try: + tags = control.list_tags_for_resource(resourceArn=arn).get("tags", {}) + except Exception: + logger.debug("Could not read tags for %s; leaving it alone", arn) + return False + return tags.get(MANAGED_TAG) == MANAGED_VALUE + + +def reap_stale_runtimes( + control: Any, + *, + max_age_minutes: int = REAP_DEFAULT_MAX_AGE_MIN, + dry_run: bool = False, + now: datetime | None = None, +) -> ReapReport: + """Delete BenchFlow-managed runtimes older than *max_age_minutes*.""" + report = ReapReport() + cutoff = (now or datetime.now(UTC)) - timedelta(minutes=max_age_minutes) + + for page in control.get_paginator("list_agent_runtimes").paginate(): + for runtime in page.get("agentRuntimes", []): + report.scanned += 1 + arn = runtime.get("agentRuntimeArn") + runtime_id = runtime.get("agentRuntimeId") + if not arn or not runtime_id: + continue + if not _is_benchflow_managed(control, arn): + report.skipped_unmanaged += 1 + continue + + created = runtime.get("createdAt") + if isinstance(created, datetime): + stamp = created if created.tzinfo else created.replace(tzinfo=UTC) + if stamp > cutoff: + report.skipped_recent += 1 + continue + + if dry_run: + report.deleted.append(runtime_id) + continue + try: + control.delete_agent_runtime(agentRuntimeId=runtime_id) + report.deleted.append(runtime_id) + logger.info("Reaped AgentCore runtime %s", runtime_id) + except Exception as exc: + report.errors.append(f"{runtime_id}: {exc}") + logger.warning( + "Failed to reap AgentCore runtime %s: %s", runtime_id, exc + ) + + return report diff --git a/tests/test_agentcore_parallel.py b/tests/test_agentcore_parallel.py new file mode 100644 index 000000000..26c6803e8 --- /dev/null +++ b/tests/test_agentcore_parallel.py @@ -0,0 +1,356 @@ +"""Parallel-safety of AgentCore image and runtime provisioning. + +A rollout is a *session*; the image and the runtime behind it are shared. That +is not an optimization — the account quotas are 5000 concurrent sessions +against 100 total runtimes with ``CreateAgentRuntime`` at 5/s, so anything +that scales with rollouts instead of with distinct images cannot run a matrix. + +These tests pin the three properties that make a fan-out safe: + +* concurrent rollouts of one task build, push, and register **once**; +* runtime identity follows the image's content, not the task name, so repeated + trials share a runtime instead of racing to create and delete one; +* ending a rollout stops its session and leaves the shared runtime alone. +""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime, timedelta +from unittest.mock import MagicMock, patch + +import pytest + +from benchflow.sandbox import agentcore_provisioning as provisioning +from benchflow.sandbox.agentcore import AgentCoreSandbox +from benchflow.task.config import SandboxConfig + + +@pytest.fixture(autouse=True) +def _clean_provisioning_cache(): + provisioning.reset_cache() + yield + provisioning.reset_cache() + + +def _make_task(tmp_path, name="demo-task", dockerfile="FROM python:3.12-slim\n"): + task_dir = tmp_path / name + task_dir.mkdir(parents=True, exist_ok=True) + (task_dir / "Dockerfile").write_text(dockerfile) + return task_dir + + +def _sandbox(task_dir, *, name="demo-task", session="run-1"): + env = AgentCoreSandbox( + environment_dir=task_dir, + environment_name=name, + session_id=session, + rollout_paths=None, + task_env_config=SandboxConfig(), + ) + env._account_id = "123456789012" + return env + + +class TestImageIdentity: + def test_identity_is_stable_across_copies_of_the_same_task(self, tmp_path): + """BenchFlow copies tasks to temp dirs; identity must survive that. + + A digest that folded in paths or mtimes would change every run and + defeat image reuse entirely. + """ + a = _make_task(tmp_path / "first") + b = _make_task(tmp_path / "second") + + assert _sandbox(a)._image_identity() == _sandbox(b)._image_identity() + + def test_identity_changes_when_the_environment_changes(self, tmp_path): + a = _make_task(tmp_path / "a") + b = _make_task(tmp_path / "b", dockerfile="FROM python:3.13-slim\n") + + assert _sandbox(a)._image_identity() != _sandbox(b)._image_identity() + + def test_identity_changes_when_a_context_file_changes(self, tmp_path): + """Skills baked into the image must produce a distinct image.""" + a = _make_task(tmp_path / "a") + b = _make_task(tmp_path / "b") + (b / "skills").mkdir() + (b / "skills" / "SKILL.md").write_text("# a skill") + + assert _sandbox(a)._image_identity() != _sandbox(b)._image_identity() + + def test_runtime_name_follows_the_image_not_the_task_name(self, tmp_path): + """Guards the delete-out-from-under-you bug in the first AgentCore cut. + + Naming the runtime after the task meant concurrent trials raced to + create one runtime and the first to finish deleted it mid-run. + """ + digest_a = "a" * 64 + digest_b = "b" * 64 + + same = provisioning.runtime_name("demo", digest_a) + again = provisioning.runtime_name("demo", digest_a) + different = provisioning.runtime_name("demo", digest_b) + + assert same == again + assert same != different + assert same[0].isalpha() + assert len(same) <= 48 + + +class TestSingleFlight: + @pytest.mark.asyncio + async def test_concurrent_rollouts_build_and_push_once(self, tmp_path): + """20 concurrent rollouts of one task must not run 20 builds.""" + task_dir = _make_task(tmp_path) + builds = {"n": 0} + + async def _build(self, image_uri, *, force_build): + builds["n"] += 1 + await asyncio.sleep(0.01) + + with ( + patch.object(AgentCoreSandbox, "_ensure_ecr_repository"), + patch.object(AgentCoreSandbox, "_ecr_image_exists", return_value=False), + patch.object(AgentCoreSandbox, "_build_and_push", new=_build), + ): + sandboxes = [_sandbox(task_dir, session=f"run-{i}") for i in range(20)] + uris = await asyncio.gather( + *(s._publish_image(force_build=False) for s in sandboxes) + ) + + assert builds["n"] == 1 + assert len(set(uris)) == 1 + + @pytest.mark.asyncio + async def test_concurrent_rollouts_register_one_runtime(self, tmp_path): + """CreateAgentRuntime is a 5/s quota — it cannot run per rollout.""" + task_dir = _make_task(tmp_path) + creates = {"n": 0} + + async def _create(self, name, image_uri): + creates["n"] += 1 + await asyncio.sleep(0.01) + return f"arn:aws:bedrock-agentcore:us-west-2:1:runtime/{name}", "rt-1" + + with patch.object(AgentCoreSandbox, "_create_or_adopt_runtime", new=_create): + sandboxes = [_sandbox(task_dir, session=f"run-{i}") for i in range(20)] + arns = await asyncio.gather( + *(s._ensure_runtime("img:tag") for s in sandboxes) + ) + + assert creates["n"] == 1 + assert len(set(arns)) == 1 + + @pytest.mark.asyncio + async def test_a_failed_build_is_not_cached(self, tmp_path): + """A transient throttle must not poison every later rollout.""" + task_dir = _make_task(tmp_path) + calls = {"n": 0} + + async def _flaky(self, image_uri, *, force_build): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("transient throttle") + + with ( + patch.object(AgentCoreSandbox, "_ensure_ecr_repository"), + patch.object(AgentCoreSandbox, "_ecr_image_exists", return_value=False), + patch.object(AgentCoreSandbox, "_build_and_push", new=_flaky), + ): + env = _sandbox(task_dir) + with pytest.raises(RuntimeError, match="transient throttle"): + await env._publish_image(force_build=False) + await env._publish_image(force_build=False) + + assert calls["n"] == 2 + + +class TestSessionTeardown: + @pytest.mark.asyncio + async def test_stop_ends_the_session_and_keeps_the_runtime(self, tmp_path): + """Deleting the runtime would tear down sibling trials still running.""" + env = _sandbox(_make_task(tmp_path)) + env.runtime_arn = "arn:aws:bedrock-agentcore:us-west-2:1:runtime/shared" + env.runtime_session_id = "s" * 40 + env._runtime_id = "shared-1" + + data = MagicMock() + control = MagicMock() + + def _client(service): + return control if service.endswith("-control") else data + + with patch.object(env, "_client", side_effect=_client): + await env.stop(delete=True) + + data.stop_runtime_session.assert_called_once() + control.delete_agent_runtime.assert_not_called() + assert env.runtime_session_id is None + + +class TestImageSizeGate: + def test_image_within_the_cap_is_accepted(self): + assert provisioning.image_size_error(1500 * 1024 * 1024, "img") is None + + def test_oversized_image_names_the_hard_quota(self): + """A 2 GB cap that is not adjustable deserves a message that says so.""" + message = provisioning.image_size_error(3000 * 1024 * 1024, "img:tag") + + assert message is not None + assert "2048" in message + assert "not" in message and "adjustable" in message + assert "daytona" in message + + +class TestReaper: + def _control(self, runtimes, tags): + control = MagicMock() + paginator = MagicMock() + paginator.paginate.return_value = [{"agentRuntimes": runtimes}] + control.get_paginator.return_value = paginator + control.list_tags_for_resource.side_effect = lambda resourceArn: { + "tags": tags.get(resourceArn, {}) + } + return control + + def test_only_benchflow_managed_runtimes_are_reaped(self): + """Never delete something another tool created in the same account.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + old = datetime.now(UTC) - timedelta(days=3) + runtimes = [ + {"agentRuntimeArn": "arn:mine", "agentRuntimeId": "mine", "createdAt": old}, + { + "agentRuntimeArn": "arn:other", + "agentRuntimeId": "other", + "createdAt": old, + }, + ] + control = self._control( + runtimes, + {"arn:mine": {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}}, + ) + + report = reap_stale_runtimes(control, max_age_minutes=60) + + assert report.deleted == ["mine"] + assert report.skipped_unmanaged == 1 + + def test_recent_runtimes_are_kept(self): + """A runtime from a run still in flight must survive cleanup.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + runtimes = [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "createdAt": datetime.now(UTC) - timedelta(minutes=5), + } + ] + control = self._control( + runtimes, + {"arn:mine": {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}}, + ) + + report = reap_stale_runtimes(control, max_age_minutes=1440) + + assert report.deleted == [] + assert report.skipped_recent == 1 + control.delete_agent_runtime.assert_not_called() + + def test_dry_run_reports_without_deleting(self): + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + runtimes = [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "createdAt": datetime.now(UTC) - timedelta(days=3), + } + ] + control = self._control( + runtimes, + {"arn:mine": {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}}, + ) + + report = reap_stale_runtimes(control, max_age_minutes=60, dry_run=True) + + assert report.deleted == ["mine"] + control.delete_agent_runtime.assert_not_called() + + def test_unreadable_tags_fail_closed(self): + """If tags can't be read, assume it isn't ours.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control( + [ + { + "agentRuntimeArn": "arn:x", + "agentRuntimeId": "x", + "createdAt": datetime.now(UTC) - timedelta(days=3), + } + ], + {}, + ) + control.list_tags_for_resource.side_effect = RuntimeError("denied") + + report = reap_stale_runtimes(control, max_age_minutes=60) + + assert report.deleted == [] + assert report.skipped_unmanaged == 1 + + +class TestCleanupCommandGating: + """``bench sandbox cleanup`` must not phone AWS unless AgentCore is set up.""" + + def test_cleanup_is_inert_without_agentcore_configuration(self, monkeypatch): + """A dev machine with AWS credentials must not trigger a live call. + + ``boto3`` ships with several unrelated extras, so importability is not + evidence that this account is being used for AgentCore runs. + """ + from benchflow.cli.sandbox import _cleanup_agentcore_runtimes + + monkeypatch.delenv("BENCHFLOW_AGENTCORE_ROLE_ARN", raising=False) + with patch("boto3.Session") as session: + assert ( + _cleanup_agentcore_runtimes(dry_run=True, max_age_minutes=60) is False + ) + session.assert_not_called() + + def test_credential_failure_degrades_instead_of_crashing(self, monkeypatch): + """Cleanup may still have Daytona work to do; don't abort the command.""" + from benchflow.cli.sandbox import _cleanup_agentcore_runtimes + + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/x") + with patch("boto3.Session", side_effect=RuntimeError("no credentials")): + assert ( + _cleanup_agentcore_runtimes(dry_run=True, max_age_minutes=60) is False + ) + + def test_daytona_failure_does_not_block_agentcore_cleanup(self, monkeypatch): + """One backend's broken credentials must not strand the other's resources. + + The Daytona path exits when DAYTONA_API_KEY is missing; before this was + isolated it aborted the whole command, silently leaving AgentCore + runtimes to accumulate against a 100-per-account quota. + """ + import typer + + from benchflow.cli import sandbox as sandbox_cli + + calls: list[str] = [] + monkeypatch.setattr(sandbox_cli, "_daytona_sdk_available", lambda: True) + monkeypatch.setattr( + sandbox_cli, + "_cleanup_agentcore_runtimes", + lambda **kw: calls.append("agentcore") or True, + ) + fake_main = MagicMock() + fake_main._cleanup_daytona_sandboxes.side_effect = typer.Exit(1) + monkeypatch.setitem(__import__("sys").modules, "benchflow.cli.main", fake_main) + + sandbox_cli.sandbox_cleanup(dry_run=True, max_age_minutes=60) + + assert calls == ["agentcore"] From b274cc431750bfb9e9721e01e422ff9eea73f86e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Mon, 27 Jul 2026 14:11:05 -0700 Subject: [PATCH 03/13] feat(agentcore): build images on CodeBuild so no local Docker is needed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AgentCore only runs images that already exist in ECR, so something must build one — and until now that was a local Docker daemon. On a machine without the resources to run containers, which is precisely when a cloud sandbox is worth reaching for, the backend was unusable. Image building is now a strategy behind `agentcore_builder`: - LocalDockerBuilder: the existing path, native and fast on an arm64 host. - CodeBuildBuilder: zips the build context to S3 and builds it on an AWS CodeBuild Graviton worker, pushing straight to ECR. Nothing is required locally — no Docker, no arm64 host, no qemu — so this also works from CI and from Windows. Default is `auto`: use Docker when a daemon actually answers, otherwise build remotely. The check probes the daemon rather than the CLI, because an installed binary with a stopped daemon is the common laptop case and finding out at build time wastes the whole provisioning path. Override with BENCHFLOW_AGENTCORE_BUILDER=docker|codebuild|auto. Verified live end to end with no `docker` on PATH at all: preflight passes, auto-selects CodeBuild, builds arm64 remotely, and the resulting session shows the file baked by the task's own Dockerfile (43s). The archive is a ZIP, not a tar.gz. CodeBuild's S3 source only unpacks ZIP; a tarball is downloaded verbatim, leaving the build directory holding the archive and failing with a misleading "Dockerfile ... no such file or directory". That cost a live build to discover, so it is pinned by a test. The 2 GB image cap is enforced on the worker too, before the push, and an oversized remote build is reported as a size problem rather than a generic build failure — otherwise the user goes hunting the wrong bug. Also measured the parallel ceiling now that runtimes are shared: 60 concurrent rollouts of one task ran on a single runtime in ~30s wall clock, 60/60 clean, no throttling, median session start ~15s. Documented alongside the quotas. --- docs/running-benchmarks.md | 28 ++ src/benchflow/sandbox/agentcore.py | 115 ++--- src/benchflow/sandbox/agentcore_builder.py | 482 +++++++++++++++++++++ tests/test_agentcore_builder.py | 175 ++++++++ tests/test_agentcore_sandbox.py | 61 ++- 5 files changed, 752 insertions(+), 109 deletions(-) create mode 100644 src/benchflow/sandbox/agentcore_builder.py create mode 100644 tests/test_agentcore_builder.py diff --git a/docs/running-benchmarks.md b/docs/running-benchmarks.md index 4546b0a23..8816ff13e 100644 --- a/docs/running-benchmarks.md +++ b/docs/running-benchmarks.md @@ -379,6 +379,29 @@ registers it as an agent runtime. The image is tagged by content digest, so a task whose environment has not changed reuses the pushed layers and the existing runtime instead of paying for a rebuild. +#### Building without Docker + +By default BenchFlow builds with a local Docker daemon when one is running, +and otherwise builds on **AWS CodeBuild** — so the backend works on a machine +with no container runtime at all, and from CI or Windows. Set a CodeBuild +service role to enable the remote path: + +```bash +export BENCHFLOW_AGENTCORE_CODEBUILD_ROLE_ARN="arn:aws:iam:::role/" +``` + +That role must be assumable by `codebuild.amazonaws.com` and able to push to +ECR, read the build bucket, and write CloudWatch logs. Your own credentials +need `codebuild:*` and S3 access. BenchFlow creates the CodeBuild project and +an `benchflow-agentcore-build--` bucket on first use (build +contexts expire after a day); override the bucket with +`BENCHFLOW_AGENTCORE_BUILD_BUCKET`. + +Builds run on a Graviton worker, so `linux/arm64` is native — no qemu. Force +either strategy with `BENCHFLOW_AGENTCORE_BUILDER=docker|codebuild|auto`. +A remote build takes roughly 40s versus roughly 20s locally on Apple silicon; +either way it happens once per distinct task image, not once per rollout. + BenchFlow also appends a small HTTP responder to the image and makes it the entrypoint. AgentCore refuses command execution for a session whose container does not answer `GET /ping` on port 8080, and task images know nothing about @@ -402,6 +425,11 @@ of the build context) and every trial and skill arm of that image opens another session against the shared runtime. Concurrency is therefore bounded by sessions, not by how many images you have. +Measured against the live service: 60 concurrent rollouts of one task ran on a +single runtime in ~30s wall clock with no throttling and no failures (median +session start ~15s). The session-creation rate quota (400/min) is the first +thing you would hit, well before the 5,000 concurrent-session ceiling. + Two quotas to check before a big run: | Quota | Default | Adjustable | diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py index e8bb3e068..bece7aed2 100644 --- a/src/benchflow/sandbox/agentcore.py +++ b/src/benchflow/sandbox/agentcore.py @@ -48,7 +48,6 @@ import io import os import shlex -import shutil import subprocess import tarfile import threading @@ -57,6 +56,7 @@ from typing import TYPE_CHECKING, Any from benchflow._paths import iter_safe_tree +from benchflow.sandbox import agentcore_builder as builders from benchflow.sandbox import agentcore_provisioning as provisioning from benchflow.sandbox._base import BaseSandbox, ExecResult, wrap_command_with_env_file from benchflow.sandbox.protocol import SandboxStartupError @@ -206,10 +206,17 @@ def preflight(cls) -> None: f"variables. Underlying error: {exc}" ) from exc - if not shutil.which("docker"): + # Docker is not required: without a local daemon, images are built on + # AWS CodeBuild instead. Only flag the case where neither route is + # usable, so the run fails now rather than at the first build. + if not builders.docker_available() and not os.environ.get( + builders.ENV_CODEBUILD_ROLE + ): raise SystemExit( - "AgentCore builds task images locally and pushes them to ECR, " - "which requires the docker CLI on PATH." + "AgentCore needs a way to build task images. Either start a " + "local Docker daemon, or set " + f"{builders.ENV_CODEBUILD_ROLE} to a CodeBuild service role so " + "images can be built in AWS without Docker." ) def _validate_definition(self) -> None: @@ -309,78 +316,24 @@ async def _publish() -> str: return await provisioning.once(cache_key, _publish) async def _build_and_push(self, image_uri: str, *, force_build: bool) -> None: - build_context, dockerfile, shim = self._materialize_build_context() - build_args = [ - "docker", - "build", - "--platform", - "linux/arm64", - "-f", - str(dockerfile), - "-t", - image_uri, - ] - if force_build: - build_args.append("--no-cache") - build_args.append(str(build_context)) - - try: - result = await asyncio.to_thread( - _run, *build_args, timeout=self.task_env_config.build_timeout_sec - ) - finally: - # The generated Dockerfile and shim only exist for the build. The - # build context is the task's own environment directory, so leaving - # them behind would litter the caller's task tree. - for path in (dockerfile, shim): - path.unlink(missing_ok=True) - if result.returncode != 0: - raise RuntimeError( - f"docker build failed (exit {result.returncode}):\n" - f"{(result.stderr or result.stdout or '').strip()[:4000]}" - ) - - await self._reject_oversized_image(image_uri) - - registry = self._ecr_registry() - await asyncio.to_thread(self._docker_login, registry) - push = await asyncio.to_thread(_run, "docker", "push", image_uri, timeout=1800) - if push.returncode != 0: - raise RuntimeError( - f"docker push to ECR failed (exit {push.returncode}):\n" - f"{(push.stderr or push.stdout or '').strip()[:4000]}" + builder = builders.select_builder( + self._client, + account_id=self._resolve_account_id(), + region=self.region, + ) + self.logger.info("Building %s via %s", image_uri, builder.name) + await builder.build_and_push( + builders.BuildRequest( + context_dir=self.environment_dir, + dockerfile_text=self._generated_dockerfile_text(), + shim_text=_PING_SHIM, + image_uri=image_uri, + registry=self._ecr_registry(), + region=self.region, + force_build=force_build, + timeout_sec=self.task_env_config.build_timeout_sec, ) - self.logger.info("Pushed AgentCore image %s", image_uri) - - async def _reject_oversized_image(self, image_uri: str) -> None: - """Fail before push when the image exceeds AgentCore's hard 2 GB cap.""" - inspect = await asyncio.to_thread( - _run, "docker", "image", "inspect", "-f", "{{.Size}}", image_uri ) - raw = (inspect.stdout or "").strip() - if inspect.returncode != 0 or not raw.isdigit(): - # Can't measure — let the push proceed rather than block on a - # diagnostic that failed for an unrelated reason. - self.logger.debug("Could not measure image size for %s", image_uri) - return - message = provisioning.image_size_error(int(raw), image_uri) - if message: - raise SandboxStartupError(message, sandbox_id=image_uri) - - def _materialize_build_context(self) -> tuple[Path, Path, Path]: - """Return (context_dir, dockerfile, shim) for an image with the shim. - - The task's own Dockerfile is left untouched on disk; the shim is - appended into a *generated* Dockerfile inside the same build context so - relative ``COPY`` paths keep resolving. Both generated files are - transient — the caller removes them once the build finishes. - """ - context = self.environment_dir - shim_path = context / provisioning.GENERATED_SHIM - shim_path.write_text(_PING_SHIM) - generated = context / provisioning.GENERATED_DOCKERFILE - generated.write_text(self._generated_dockerfile_text()) - return context, generated, shim_path def _generated_dockerfile_text(self) -> str: """The Dockerfile BenchFlow actually builds, as text. @@ -426,20 +379,6 @@ def _ecr_image_exists(self, tag: str) -> bool: except ClientError: return False - def _docker_login(self, registry: str) -> None: - token = self._client("ecr").get_authorization_token() - blob = token["authorizationData"][0]["authorizationToken"] - _user, password = base64.b64decode(blob).decode().split(":", 1) - proc = subprocess.run( - ["docker", "login", "--username", "AWS", "--password-stdin", registry], - input=password, - capture_output=True, - text=True, - timeout=120, - ) - if proc.returncode != 0: - raise RuntimeError(f"ECR docker login failed: {(proc.stderr or '')[:1000]}") - async def _ensure_runtime(self, image_uri: str) -> str: """Resolve the shared agent runtime for *image_uri*, creating it once. diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py new file mode 100644 index 000000000..9ed30f742 --- /dev/null +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -0,0 +1,482 @@ +"""How a task image gets built and pushed to ECR for the AgentCore backend. + +AgentCore only accepts an image that already exists in ECR, so something has +to build one. Two strategies: + +* :class:`LocalDockerBuilder` shells out to the local Docker daemon. Fast when + the host is arm64 (a native build), and the obvious choice on a developer + machine. +* :class:`CodeBuildBuilder` ships the build context to S3 and builds it in AWS + CodeBuild on a Graviton worker, pushing straight to ECR. Nothing is required + locally — no Docker, no arm64 host, no qemu — which is what makes the + backend usable from a laptop without a container runtime, from CI, and from + Windows. + +The default is ``auto``: use Docker when a working daemon is present, and fall +back to CodeBuild when it is not. Selection is explicit in the logs, because +"why did this build take four minutes" should never be a mystery. +""" + +from __future__ import annotations + +import asyncio +import io +import json +import logging +import os +import subprocess +import time +import uuid +import zipfile +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Protocol + +from benchflow.sandbox import agentcore_provisioning as provisioning + +logger = logging.getLogger("benchflow").getChild("agentcore-builder") + +ENV_BUILDER = "BENCHFLOW_AGENTCORE_BUILDER" +ENV_CODEBUILD_ROLE = "BENCHFLOW_AGENTCORE_CODEBUILD_ROLE_ARN" +ENV_BUILD_BUCKET = "BENCHFLOW_AGENTCORE_BUILD_BUCKET" + +CODEBUILD_PROJECT = "benchflow-agentcore-builder" +# Graviton worker: builds linux/arm64 natively, so no qemu emulation. +CODEBUILD_IMAGE = "aws/codebuild/amazonlinux-aarch64-standard:3.0" +CODEBUILD_COMPUTE = "BUILD_GENERAL1_LARGE" +_CODEBUILD_POLL_SEC = 10 +_CODEBUILD_TIMEOUT_MIN = 60 + + +@dataclass(frozen=True) +class BuildRequest: + """Everything needed to produce one image in ECR.""" + + context_dir: Path + dockerfile_text: str + shim_text: str + image_uri: str + registry: str + region: str + force_build: bool + timeout_sec: float | None + + +class Builder(Protocol): + name: str + + async def build_and_push(self, request: BuildRequest) -> None: ... + + +@contextmanager +def materialized(request: BuildRequest) -> Iterator[Path]: + """Write the generated Dockerfile and shim, yield the Dockerfile, clean up. + + The build context is the task's own environment directory, so these files + must not survive the build — otherwise every run would litter the caller's + task tree. + """ + context = request.context_dir + shim = context / provisioning.GENERATED_SHIM + dockerfile = context / provisioning.GENERATED_DOCKERFILE + shim.write_text(request.shim_text) + dockerfile.write_text(request.dockerfile_text) + try: + yield dockerfile + finally: + for path in (dockerfile, shim): + path.unlink(missing_ok=True) + + +def _run(*args: str, timeout: float | None = None, input_text: str | None = None): + return subprocess.run( + args, capture_output=True, text=True, timeout=timeout, input=input_text + ) + + +def docker_available() -> bool: + """True when a Docker daemon is actually reachable. + + Checks the daemon, not just the CLI: an installed ``docker`` binary with a + stopped daemon is the common laptop case, and discovering that only at + build time would waste the whole image push. + """ + import shutil + + if not shutil.which("docker"): + return False + try: + return ( + _run( + "docker", "info", "--format", "{{.ServerVersion}}", timeout=15 + ).returncode + == 0 + ) + except (OSError, subprocess.SubprocessError): + return False + + +class LocalDockerBuilder: + """Build with the local Docker daemon.""" + + name = "docker" + + def __init__(self, client_factory: Any) -> None: + self._client = client_factory + + async def build_and_push(self, request: BuildRequest) -> None: + with materialized(request) as dockerfile: + args = [ + "docker", + "build", + "--platform", + "linux/arm64", + "-f", + str(dockerfile), + "-t", + request.image_uri, + ] + if request.force_build: + args.append("--no-cache") + args.append(str(request.context_dir)) + result = await asyncio.to_thread(_run, *args, timeout=request.timeout_sec) + if result.returncode != 0: + raise RuntimeError( + f"docker build failed (exit {result.returncode}):\n" + f"{(result.stderr or result.stdout or '').strip()[:4000]}" + ) + + await self._reject_oversized(request.image_uri) + await asyncio.to_thread(self._login, request.registry) + push = await asyncio.to_thread( + _run, "docker", "push", request.image_uri, timeout=1800 + ) + if push.returncode != 0: + raise RuntimeError( + f"docker push to ECR failed (exit {push.returncode}):\n" + f"{(push.stderr or push.stdout or '').strip()[:4000]}" + ) + logger.info("Pushed AgentCore image %s (local docker)", request.image_uri) + + async def _reject_oversized(self, image_uri: str) -> None: + inspect = await asyncio.to_thread( + _run, "docker", "image", "inspect", "-f", "{{.Size}}", image_uri + ) + raw = (inspect.stdout or "").strip() + if inspect.returncode != 0 or not raw.isdigit(): + logger.debug("Could not measure image size for %s", image_uri) + return + message = provisioning.image_size_error(int(raw), image_uri) + if message: + from benchflow.sandbox.protocol import SandboxStartupError + + raise SandboxStartupError(message, sandbox_id=image_uri) + + def _login(self, registry: str) -> None: + import base64 + + token = self._client("ecr").get_authorization_token() + blob = token["authorizationData"][0]["authorizationToken"] + _user, password = base64.b64decode(blob).decode().split(":", 1) + proc = _run( + "docker", + "login", + "--username", + "AWS", + "--password-stdin", + registry, + timeout=120, + input_text=password, + ) + if proc.returncode != 0: + raise RuntimeError(f"ECR docker login failed: {(proc.stderr or '')[:1000]}") + + +# Runs on the CodeBuild worker. The size gate lives here too so an oversized +# image is rejected before the push rather than surfacing later as an opaque +# AgentCore runtime error. +_BUILDSPEC = { + "version": "0.2", + "phases": { + "pre_build": { + "commands": [ + "aws ecr get-login-password --region $AWS_REGION " + "| docker login --username AWS --password-stdin $BF_REGISTRY", + ] + }, + "build": { + "commands": [ + f"docker build --platform linux/arm64 " + f"-f {provisioning.GENERATED_DOCKERFILE} -t $BF_IMAGE_URI .", + 'SIZE=$(docker image inspect -f "{{.Size}}" $BF_IMAGE_URI)', + "MB=$((SIZE / 1048576))", + 'echo "benchflow: image size ${MB} MB"', + f'if [ "$MB" -gt {provisioning.MAX_IMAGE_MB} ]; then ' + f'echo "BENCHFLOW_IMAGE_TOO_LARGE ${{MB}}"; exit 1; fi', + "docker push $BF_IMAGE_URI", + ] + }, + }, +} + + +class CodeBuildBuilder: + """Build on AWS CodeBuild (Graviton), requiring nothing locally.""" + + name = "codebuild" + + def __init__(self, client_factory: Any, account_id: str, region: str) -> None: + self._client = client_factory + self._account_id = account_id + self._region = region + + @property + def _bucket(self) -> str: + return ( + os.environ.get(ENV_BUILD_BUCKET) + or f"benchflow-agentcore-build-{self._account_id}-{self._region}" + ) + + def _role_arn(self) -> str: + role = os.environ.get(ENV_CODEBUILD_ROLE) + if role: + return role + raise RuntimeError( + "Remote image builds need a CodeBuild service role. Set " + f"{ENV_CODEBUILD_ROLE} to a role assumable by codebuild.amazonaws.com " + "that can push to ECR, read the build bucket, and write CloudWatch " + "logs. (Alternatively install Docker locally and set " + f"{ENV_BUILDER}=docker.)" + ) + + async def build_and_push(self, request: BuildRequest) -> None: + key = f"contexts/{uuid.uuid4().hex}.zip" + archive = await asyncio.to_thread(self._package, request) + await asyncio.to_thread(self._ensure_bucket) + await asyncio.to_thread( + self._client("s3").put_object, + Bucket=self._bucket, + Key=key, + Body=archive, + ) + try: + await asyncio.to_thread(self._ensure_project) + await self._run_build(request, key) + finally: + try: + await asyncio.to_thread( + self._client("s3").delete_object, Bucket=self._bucket, Key=key + ) + except Exception as exc: + logger.debug("Could not delete build context %s: %s", key, exc) + + def _package(self, request: BuildRequest) -> bytes: + """Zip the build context with the generated Dockerfile and shim inside. + + ZIP, not tar.gz: CodeBuild's S3 source type only unpacks ZIP archives. + A tarball is downloaded verbatim, so the build directory ends up + holding the archive itself and the build fails with a misleading + ``Dockerfile ... no such file or directory``. + """ + buffer = io.BytesIO() + with ( + materialized(request), + zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive, + ): + for path in sorted(request.context_dir.rglob("*")): + if path.is_symlink() or not path.is_file(): + # Symlinks are skipped for the same reason as on upload + # (#411): they must not pull host files into the image. + continue + archive.write( + path, arcname=path.relative_to(request.context_dir).as_posix() + ) + return buffer.getvalue() + + def _ensure_bucket(self) -> None: + from botocore.exceptions import ClientError + + s3 = self._client("s3") + try: + s3.head_bucket(Bucket=self._bucket) + return + except ClientError: + pass + kwargs: dict[str, Any] = {"Bucket": self._bucket} + if self._region != "us-east-1": + kwargs["CreateBucketConfiguration"] = {"LocationConstraint": self._region} + try: + s3.create_bucket(**kwargs) + except ClientError as exc: + if exc.response["Error"]["Code"] not in { + "BucketAlreadyOwnedByYou", + "BucketAlreadyExists", + }: + raise + for call, kw in ( + ( + s3.put_public_access_block, + { + "Bucket": self._bucket, + "PublicAccessBlockConfiguration": { + "BlockPublicAcls": True, + "IgnorePublicAcls": True, + "BlockPublicPolicy": True, + "RestrictPublicBuckets": True, + }, + }, + ), + ( + s3.put_bucket_lifecycle_configuration, + { + "Bucket": self._bucket, + # Build contexts are consumed within minutes; expiring them + # keeps a benchmark run from accruing storage cost forever. + "LifecycleConfiguration": { + "Rules": [ + { + "ID": "benchflow-expire-build-contexts", + "Status": "Enabled", + "Filter": {"Prefix": "contexts/"}, + "Expiration": {"Days": 1}, + } + ] + }, + }, + ), + ): + try: + call(**kw) + except Exception as exc: + logger.debug("Bucket hardening step skipped: %s", exc) + + def _ensure_project(self) -> None: + from botocore.exceptions import ClientError + + codebuild = self._client("codebuild") + config = { + "source": {"type": "S3", "location": f"{self._bucket}/placeholder"}, + "artifacts": {"type": "NO_ARTIFACTS"}, + "environment": { + "type": "ARM_CONTAINER", + "image": CODEBUILD_IMAGE, + "computeType": CODEBUILD_COMPUTE, + # Required to run a Docker daemon inside the build container. + "privilegedMode": True, + }, + "serviceRole": self._role_arn(), + } + try: + codebuild.create_project(name=CODEBUILD_PROJECT, **config) + logger.info("Created CodeBuild project %s", CODEBUILD_PROJECT) + except ClientError as exc: + if exc.response["Error"]["Code"] != "ResourceAlreadyExistsException": + raise + + async def _run_build(self, request: BuildRequest, key: str) -> None: + codebuild = self._client("codebuild") + started = await asyncio.to_thread( + codebuild.start_build, + projectName=CODEBUILD_PROJECT, + sourceTypeOverride="S3", + sourceLocationOverride=f"{self._bucket}/{key}", + buildspecOverride=json.dumps(_BUILDSPEC), + environmentVariablesOverride=[ + {"name": "BF_IMAGE_URI", "value": request.image_uri}, + {"name": "BF_REGISTRY", "value": request.registry}, + ], + timeoutInMinutesOverride=_CODEBUILD_TIMEOUT_MIN, + ) + build_id = started["build"]["id"] + logger.info( + "CodeBuild %s building %s (arm64, no local docker)", + build_id, + request.image_uri, + ) + + deadline = time.monotonic() + ( + request.timeout_sec or _CODEBUILD_TIMEOUT_MIN * 60 + ) + while time.monotonic() < deadline: + await asyncio.sleep(_CODEBUILD_POLL_SEC) + builds = await asyncio.to_thread(codebuild.batch_get_builds, ids=[build_id]) + build = builds["builds"][0] + if build["buildStatus"] == "IN_PROGRESS": + continue + if build["buildStatus"] == "SUCCEEDED": + logger.info("Pushed AgentCore image %s (codebuild)", request.image_uri) + return + raise self._build_failure(build, request.image_uri) + raise TimeoutError( + f"CodeBuild {build_id} did not finish within the build timeout" + ) + + def _build_failure(self, build: dict[str, Any], image_uri: str) -> Exception: + """Turn a failed build into the most specific error we can offer.""" + from benchflow.sandbox.protocol import SandboxStartupError + + phases = build.get("phases") or [] + detail = "" + for phase in phases: + for context in phase.get("contexts") or []: + message = context.get("message") + if message: + detail = message + tail = self._log_tail(build) + if "BENCHFLOW_IMAGE_TOO_LARGE" in tail: + return SandboxStartupError( + provisioning.image_size_error( + (provisioning.MAX_IMAGE_MB + 1) * 1024 * 1024, image_uri + ) + or "image too large", + sandbox_id=image_uri, + ) + return RuntimeError( + f"CodeBuild {build['id']} failed ({build['buildStatus']}): " + f"{detail or 'see CloudWatch logs'}\n{tail[-2000:]}" + ) + + def _log_tail(self, build: dict[str, Any]) -> str: + logs = build.get("logs") or {} + group, stream = logs.get("groupName"), logs.get("streamName") + if not group or not stream: + return "" + try: + events = self._client("logs").get_log_events( + logGroupName=group, logStreamName=stream, limit=100 + ) + return "".join(e.get("message", "") for e in events.get("events", [])) + except Exception: + return "" + + +def select_builder( + client_factory: Any, + *, + account_id: str, + region: str, + preference: str | None = None, +) -> Builder: + """Choose a build strategy. + + ``auto`` (the default) prefers a working local Docker daemon and otherwise + builds remotely, so the backend works on a machine with no container + runtime without the caller configuring anything extra. + """ + choice = (preference or os.environ.get(ENV_BUILDER) or "auto").strip().lower() + if choice not in {"auto", "docker", "codebuild"}: + raise ValueError( + f"Invalid {ENV_BUILDER}={choice!r}; use auto, docker, or codebuild." + ) + + if choice == "docker": + return LocalDockerBuilder(client_factory) + if choice == "codebuild": + return CodeBuildBuilder(client_factory, account_id, region) + + if docker_available(): + return LocalDockerBuilder(client_factory) + logger.info("No local Docker daemon; building images on AWS CodeBuild") + return CodeBuildBuilder(client_factory, account_id, region) diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py new file mode 100644 index 000000000..d3ec90540 --- /dev/null +++ b/tests/test_agentcore_builder.py @@ -0,0 +1,175 @@ +"""Image build strategy selection for the AgentCore backend. + +AgentCore only runs images that already exist in ECR, so something must build +one. Requiring a local Docker daemon defeats the point of a cloud backend on a +machine that has no room to run containers, so the builder falls back to AWS +CodeBuild on a Graviton worker when no daemon is reachable. +""" + +from __future__ import annotations + +import zipfile +from io import BytesIO +from unittest.mock import MagicMock, patch + +import pytest + +from benchflow.sandbox import agentcore_builder as builders + + +def _request(tmp_path, **overrides): + (tmp_path / "Dockerfile").write_text("FROM python:3.12-slim\n") + defaults = dict( + context_dir=tmp_path, + dockerfile_text="FROM python:3.12-slim\n", + shim_text="print('shim')\n", + image_uri="1.dkr.ecr.us-west-2.amazonaws.com/repo:tag", + registry="1.dkr.ecr.us-west-2.amazonaws.com", + region="us-west-2", + force_build=False, + timeout_sec=None, + ) + defaults.update(overrides) + return builders.BuildRequest(**defaults) + + +class TestBuilderSelection: + def test_prefers_local_docker_when_the_daemon_is_up(self): + with patch.object(builders, "docker_available", return_value=True): + builder = builders.select_builder( + MagicMock(), account_id="1", region="us-west-2" + ) + + assert isinstance(builder, builders.LocalDockerBuilder) + + def test_falls_back_to_codebuild_without_a_daemon(self): + """The whole point: usable on a machine with no container runtime.""" + with patch.object(builders, "docker_available", return_value=False): + builder = builders.select_builder( + MagicMock(), account_id="1", region="us-west-2" + ) + + assert isinstance(builder, builders.CodeBuildBuilder) + + def test_explicit_preference_overrides_detection(self): + with patch.object(builders, "docker_available", return_value=True): + builder = builders.select_builder( + MagicMock(), account_id="1", region="us-west-2", preference="codebuild" + ) + + assert isinstance(builder, builders.CodeBuildBuilder) + + def test_invalid_preference_is_rejected(self): + with pytest.raises(ValueError, match="auto, docker, or codebuild"): + builders.select_builder( + MagicMock(), account_id="1", region="us-west-2", preference="podman" + ) + + def test_a_stopped_daemon_counts_as_unavailable(self): + """An installed CLI with a dead daemon is the common laptop case. + + Discovering that at build time would waste the whole provisioning path. + """ + with ( + patch("shutil.which", return_value="/usr/bin/docker"), + patch.object(builders, "_run", return_value=MagicMock(returncode=1)), + ): + assert builders.docker_available() is False + + def test_codebuild_without_a_role_says_what_to_set(self, monkeypatch): + monkeypatch.delenv(builders.ENV_CODEBUILD_ROLE, raising=False) + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + + with pytest.raises(RuntimeError) as excinfo: + builder._role_arn() + + assert builders.ENV_CODEBUILD_ROLE in str(excinfo.value) + assert "codebuild.amazonaws.com" in str(excinfo.value) + + +class TestCodeBuildPackaging: + def test_archive_carries_the_generated_dockerfile_and_shim(self, tmp_path): + """CodeBuild only sees the archive, so the scaffolding must be inside.""" + (tmp_path / "data.txt").write_text("payload") + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + + blob = builder._package(_request(tmp_path)) + + with zipfile.ZipFile(BytesIO(blob)) as archive: + names = set(archive.namelist()) + from benchflow.sandbox import agentcore_provisioning as provisioning + + assert provisioning.GENERATED_DOCKERFILE in names + assert provisioning.GENERATED_SHIM in names + assert "Dockerfile" in names + assert "data.txt" in names + + def test_packaging_leaves_no_scaffolding_behind(self, tmp_path): + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + request = _request(tmp_path) + before = {p.name for p in tmp_path.iterdir()} + + builder._package(request) + + assert {p.name for p in tmp_path.iterdir()} == before + + def test_symlinks_are_not_packaged(self, tmp_path): + """Guards #411 — a task symlink must not ship host files to AWS.""" + secret = tmp_path.parent / "host-secret.txt" + secret.write_text("do not upload me") + (tmp_path / "link.txt").symlink_to(secret) + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + + blob = builder._package(_request(tmp_path)) + + with zipfile.ZipFile(BytesIO(blob)) as archive: + assert "link.txt" not in set(archive.namelist()) + + def test_buildspec_enforces_the_image_size_cap_remotely(self): + """The 2 GB cap must be caught on the worker, before the push.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + commands = " ".join(builders._BUILDSPEC["phases"]["build"]["commands"]) + + assert str(provisioning.MAX_IMAGE_MB) in commands + assert "BENCHFLOW_IMAGE_TOO_LARGE" in commands + # The push must come after the gate, or an oversized image still lands. + assert commands.index("BENCHFLOW_IMAGE_TOO_LARGE") < commands.index( + "docker push" + ) + + def test_buildspec_builds_arm64(self): + """AgentCore runs arm64 only; an x86 image would fail opaquely.""" + commands = " ".join(builders._BUILDSPEC["phases"]["build"]["commands"]) + + assert "--platform linux/arm64" in commands + + def test_oversized_remote_build_is_reported_as_a_size_problem(self): + """A generic 'build failed' would send the user hunting the wrong bug.""" + from benchflow.sandbox.protocol import SandboxStartupError + + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + build = {"id": "b:1", "buildStatus": "FAILED", "phases": []} + + with patch.object( + builders.CodeBuildBuilder, + "_log_tail", + return_value="BENCHFLOW_IMAGE_TOO_LARGE 3200", + ): + error = builder._build_failure(build, "repo:tag") + + assert isinstance(error, SandboxStartupError) + assert "2048" in str(error) + + def test_context_is_zipped_not_tarred(self, tmp_path): + """CodeBuild's S3 source only unpacks ZIP. + + A tar.gz is downloaded verbatim, leaving the build directory holding + the archive itself — which surfaces as a confusing "Dockerfile ... no + such file or directory" from docker build. + """ + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + + blob = builder._package(_request(tmp_path)) + + assert zipfile.is_zipfile(BytesIO(blob)) diff --git a/tests/test_agentcore_sandbox.py b/tests/test_agentcore_sandbox.py index d4a6edc9c..2c46a342d 100644 --- a/tests/test_agentcore_sandbox.py +++ b/tests/test_agentcore_sandbox.py @@ -19,6 +19,7 @@ from __future__ import annotations import base64 +import contextlib import os import tarfile from unittest.mock import MagicMock, patch @@ -302,38 +303,57 @@ async def test_oversized_inline_write_is_refused(self, sandbox): class TestImagePreparation: - def test_generated_dockerfile_adds_the_ping_shim(self, sandbox, tmp_path): + def _request(self, sandbox): + from benchflow.sandbox import agentcore_builder as builders + from benchflow.sandbox.agentcore import _PING_SHIM + + return builders.BuildRequest( + context_dir=sandbox.environment_dir, + dockerfile_text=sandbox._generated_dockerfile_text(), + shim_text=_PING_SHIM, + image_uri="reg/repo:tag", + registry="reg", + region="us-west-2", + force_build=False, + timeout_sec=None, + ) + + def test_generated_dockerfile_adds_the_ping_shim(self, sandbox): """Without a /ping responder the microVM 500s on every command.""" - context, dockerfile, shim = sandbox._materialize_build_context() - text = dockerfile.read_text() + text = sandbox._generated_dockerfile_text() assert "FROM python:3.12-slim" in text assert "benchflow_agentcore_shim.py" in text assert "EXPOSE 8080" in text - assert shim.exists() - assert shim.parent == context def test_task_dockerfile_is_not_modified_in_place(self, sandbox, tmp_path): - original = (tmp_path / "Dockerfile").read_text() - sandbox._materialize_build_context() + from benchflow.sandbox import agentcore_builder as builders - assert (tmp_path / "Dockerfile").read_text() == original + original = (tmp_path / "Dockerfile").read_text() + with builders.materialized(self._request(sandbox)) as dockerfile: + assert dockerfile.exists() + assert (tmp_path / "Dockerfile").read_text() == original - @pytest.mark.asyncio - async def test_build_leaves_no_generated_files_in_the_task_tree( - self, sandbox, tmp_path - ): + def test_generated_files_never_survive_the_build(self, sandbox, tmp_path): """The build context is the task's own directory; don't litter it.""" + from benchflow.sandbox import agentcore_builder as builders + before = {p.name for p in tmp_path.iterdir()} + with builders.materialized(self._request(sandbox)): + assert {p.name for p in tmp_path.iterdir()} != before - async def _to_thread(fn, *args, **kwargs): - return MagicMock(returncode=0, stdout="sha256:" + "a" * 64, stderr="") + assert {p.name for p in tmp_path.iterdir()} == before - with ( - patch("benchflow.sandbox.agentcore.asyncio.to_thread", new=_to_thread), - patch.object(sandbox, "_ecr_registry", return_value="reg"), - ): - await sandbox._publish_image(force_build=False) + def test_generated_files_are_cleaned_up_after_a_failed_build( + self, sandbox, tmp_path + ): + """A build that raises must not leave scaffolding behind either.""" + from benchflow.sandbox import agentcore_builder as builders + + before = {p.name for p in tmp_path.iterdir()} + with contextlib.suppress(RuntimeError): # noqa: SIM117 + with builders.materialized(self._request(sandbox)): + raise RuntimeError("build blew up") assert {p.name for p in tmp_path.iterdir()} == before @@ -345,9 +365,8 @@ def test_docker_image_config_is_honoured(self, tmp_path): rollout_paths=None, task_env_config=SandboxConfig(docker_image="python:3.12-slim"), ) - _context, dockerfile, _shim = env._materialize_build_context() - assert "FROM python:3.12-slim" in dockerfile.read_text() + assert "FROM python:3.12-slim" in env._generated_dockerfile_text() class TestCapabilityGating: From 935f67f66f8508925ccc4fae31f045480d466436 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Mon, 27 Jul 2026 17:04:30 -0700 Subject: [PATCH 04/13] fix(agentcore): address P1 review findings on cleanup, context, and adoption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All five blocking findings from the PR #937 review, each re-verified against the live service rather than against invented response shapes. 1. Cleanup selected fresh runtimes (fail-open). `ListAgentRuntimes` returns only `lastUpdatedAt` — there is no `createdAt` in the list shape, though `GetAgentRuntime` has both. Reading the wrong field yielded None for every runtime, which skipped the age comparison entirely, so a one-day policy selected minutes-old runtimes. Age resolution now accepts either field and an unknown age is treated as not-stale. Verified live: the reviewer's reproduction now reports deleted=0, recent=1. 2. CodeBuild ignored `.dockerignore`. The local daemon honors it natively while the remote path zipped every regular file, uploading ignored files — including secrets — to S3, and letting an ignored file perturb the image digest. Both packaging and identity now walk one canonical context (`iter_context_files`) that applies `.dockerignore` (comments, `!` re-includes, directory prefixes), with the generated scaffolding added explicitly to the archive so the two views cannot drift. 3. Stale images could be adopted. The conflict path fetched the existing image URI and discarded it. It now compares, updates the runtime when it differs, and re-reads the runtime afterwards to fail closed if it still does not match. Runtimes are also bound to the immutable `repo@sha256:...` digest resolved from ECR rather than to a mutable tag, so what a runtime contains stays knowable after a rebuild or repository change. 4. Compose tasks were accepted but only one container was built, so the agent ran without its side services and the resulting failure was scored against the agent. `supports_compose` now lives on the provider registry; the capability gate refuses multi-service tasks during planning and the sandbox refuses them at construction. Docker and Daytona-DinD are unaffected. 5. The symlink regression test asserted against base64-encoded shell commands, where the plaintext could never appear. It now decodes the staged chunks and inspects the real tar, with a positive control. Mutation-checked: reverting the symlink skip makes it fail, which it did not before. --- src/benchflow/sandbox/_compose.py | 22 ++ src/benchflow/sandbox/agentcore.py | 103 +++++++-- src/benchflow/sandbox/agentcore_builder.py | 18 +- .../sandbox/agentcore_provisioning.py | 67 +++++- src/benchflow/sandbox/agentcore_reaper.py | 33 ++- src/benchflow/sandbox/providers.py | 16 +- src/benchflow/task/runtime_capabilities.py | 31 +++ tests/test_agentcore_builder.py | 63 +++++ tests/test_agentcore_parallel.py | 218 ++++++++++++++++++ tests/test_agentcore_sandbox.py | 33 ++- 10 files changed, 562 insertions(+), 42 deletions(-) diff --git a/src/benchflow/sandbox/_compose.py b/src/benchflow/sandbox/_compose.py index 3cec98cfd..2217aacff 100644 --- a/src/benchflow/sandbox/_compose.py +++ b/src/benchflow/sandbox/_compose.py @@ -28,6 +28,28 @@ ) +#: Filenames Docker Compose recognizes for a task-supplied topology. +COMPOSE_DEFINITION_NAMES = ( + "docker-compose.yaml", + "docker-compose.yml", + "compose.yaml", + "compose.yml", +) + + +def compose_definition_path(environment_dir: Path) -> Path | None: + """The task's own compose file in *environment_dir*, if it has one. + + Single-container backends use this to refuse a multi-service task up front + rather than silently building only the agent's container. + """ + for name in COMPOSE_DEFINITION_NAMES: + candidate = environment_dir / name + if candidate.is_file(): + return candidate + return None + + def is_compose_up_network_race_error(message: str) -> bool: """Return whether *message* is a retryable compose-up network create race.""" return bool(_COMPOSE_UP_NETWORK_RACE_ERROR.search(message)) diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py index bece7aed2..e70f24516 100644 --- a/src/benchflow/sandbox/agentcore.py +++ b/src/benchflow/sandbox/agentcore.py @@ -48,7 +48,6 @@ import io import os import shlex -import subprocess import tarfile import threading import uuid @@ -59,6 +58,7 @@ from benchflow.sandbox import agentcore_builder as builders from benchflow.sandbox import agentcore_provisioning as provisioning from benchflow.sandbox._base import BaseSandbox, ExecResult, wrap_command_with_env_file +from benchflow.sandbox._compose import compose_definition_path from benchflow.sandbox.protocol import SandboxStartupError from benchflow.task.config import SandboxConfig from benchflow.task.paths import RolloutPaths @@ -132,12 +132,6 @@ def log_message(self, *_args): ''' -def _run(*args: str, timeout: float | None = None, cwd: str | None = None): - return subprocess.run( - args, capture_output=True, text=True, timeout=timeout, cwd=cwd - ) - - class AgentCoreSandbox(BaseSandbox): """Sandbox backend for Bedrock AgentCore Runtime microVMs.""" @@ -226,6 +220,18 @@ def _validate_definition(self) -> None: f"No Dockerfile found in {self.environment_dir} and no " "docker_image specified in task config." ) + compose = compose_definition_path(self.environment_dir) + if compose is not None: + # AgentCore builds and runs exactly one container. Accepting a + # compose task would launch the agent's container without its side + # services, and the resulting failure would be scored against the + # agent rather than reported as an unsupported environment. + raise ValueError( + f"{compose.name} found in {self.environment_dir}: AgentCore is a " + "single-container backend and cannot start compose side " + "services. Run multi-service tasks on the docker or daytona " + "sandbox." + ) # ------------------------------------------------------------- lifecycle @@ -295,26 +301,47 @@ def _image_identity(self) -> tuple[str, str]: return digest, provisioning.image_tag(self.environment_name, digest) async def _publish_image(self, *, force_build: bool) -> str: - """Build the task image with the ping shim and push it to ECR. + """Build and push the task image, returning its **immutable** URI. Deduplicated process-wide by content digest: a fan-out of rollouts over one task builds and pushes exactly once, and every later run of an unchanged task skips the build entirely on the ECR existence check. + + The value returned is ``repo@sha256:...``, not ``repo:tag``. A tag is + mutable — ``force_build``, a rebased base image, or a repository change + can all move it — so binding a runtime to a tag makes the runtime's + contents unknowable after the fact. The registry digest pins exactly + the bytes that were pushed. """ _digest, tag = self._image_identity() - image_uri = f"{self._ecr_registry()}/{self._ecr_repository}:{tag}" - cache_key = f"image:{image_uri}:{force_build}" + tagged_uri = f"{self._ecr_registry()}/{self._ecr_repository}:{tag}" + cache_key = f"image:{tagged_uri}:{force_build}" async def _publish() -> str: await asyncio.to_thread(self._ensure_ecr_repository) - if not force_build and await asyncio.to_thread(self._ecr_image_exists, tag): - self.logger.info("Reusing published AgentCore image %s", image_uri) - return image_uri - await self._build_and_push(image_uri, force_build=force_build) - return image_uri + if force_build or not await asyncio.to_thread(self._ecr_image_exists, tag): + await self._build_and_push(tagged_uri, force_build=force_build) + else: + self.logger.info("Reusing published AgentCore image %s", tagged_uri) + return await asyncio.to_thread(self._resolve_image_digest, tag) return await provisioning.once(cache_key, _publish) + def _resolve_image_digest(self, tag: str) -> str: + """Resolve *tag* to the immutable ``repo@sha256:...`` reference.""" + response = self._client("ecr").describe_images( + repositoryName=self._ecr_repository, imageIds=[{"imageTag": tag}] + ) + details = response.get("imageDetails") or [] + if not details or not details[0].get("imageDigest"): + raise RuntimeError( + f"ECR did not report a digest for {self._ecr_repository}:{tag}; " + "cannot bind an AgentCore runtime to a verifiable image." + ) + return ( + f"{self._ecr_registry()}/{self._ecr_repository}@{details[0]['imageDigest']}" + ) + async def _build_and_push(self, image_uri: str, *, force_build: bool) -> None: builder = builders.select_builder( self._client, @@ -437,8 +464,33 @@ async def _create_or_adopt_runtime( ) if existing is None: raise - arn, runtime_id, _image = existing + arn, runtime_id, current_image = existing + if current_image != image_uri: + # Adopting a runtime bound to different bytes would silently run + # the wrong environment. The name encodes the build-context + # digest, but that is not the same as the pushed image: a + # repository change or a force rebuild can move what the name + # resolves to. Re-point the runtime and wait for it to settle. + self.logger.warning( + "AgentCore runtime %s points at %s, not %s — updating", + arn, + current_image, + image_uri, + ) + await asyncio.to_thread( + control.update_agent_runtime, + agentRuntimeId=runtime_id, + agentRuntimeArtifact={ + "containerConfiguration": {"containerUri": image_uri} + }, + roleArn=self._require_role_arn(), + networkConfiguration={"networkMode": "PUBLIC"}, + protocolConfiguration=_PROTOCOL_CONFIGURATION, + ) await asyncio.to_thread(self._wait_ready, control, runtime_id) + await asyncio.to_thread( + self._verify_runtime_image, control, runtime_id, image_uri + ) self.logger.info("Adopted existing AgentCore runtime %s", arn) return arn, runtime_id @@ -451,6 +503,25 @@ async def _create_or_adopt_runtime( ) return created["agentRuntimeArn"], runtime_id + @staticmethod + def _verify_runtime_image(control: Any, runtime_id: str, image_uri: str) -> None: + """Fail closed if the runtime is not bound to the image we expect. + + The adoption path is the one place a runtime can predate this process, + so it is the one place a mismatch could otherwise go unnoticed and run + an agent against the wrong environment. + """ + detail = control.get_agent_runtime(agentRuntimeId=runtime_id) + artifact = detail.get("agentRuntimeArtifact") or {} + bound = (artifact.get("containerConfiguration") or {}).get("containerUri") + if bound != image_uri: + raise SandboxStartupError( + f"AgentCore runtime {runtime_id} is bound to {bound!r} but this " + f"task needs {image_uri!r}. Refusing to run against the wrong " + "image.", + sandbox_id=runtime_id, + ) + def _lifecycle_configuration(self) -> dict[str, int]: """Session idle/lifetime caps. diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py index 9ed30f742..f05cc0fb6 100644 --- a/src/benchflow/sandbox/agentcore_builder.py +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -285,14 +285,16 @@ def _package(self, request: BuildRequest) -> bytes: materialized(request), zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive, ): - for path in sorted(request.context_dir.rglob("*")): - if path.is_symlink() or not path.is_file(): - # Symlinks are skipped for the same reason as on upload - # (#411): they must not pull host files into the image. - continue - archive.write( - path, arcname=path.relative_to(request.context_dir).as_posix() - ) + for path, relative in provisioning.iter_context_files(request.context_dir): + archive.write(path, arcname=relative) + # The generated scaffolding is excluded from the canonical walk + # (it must not affect image identity), so add it explicitly — + # CodeBuild only ever sees this archive. + for generated in ( + provisioning.GENERATED_DOCKERFILE, + provisioning.GENERATED_SHIM, + ): + archive.write(request.context_dir / generated, arcname=generated) return buffer.getvalue() def _ensure_bucket(self) -> None: diff --git a/src/benchflow/sandbox/agentcore_provisioning.py b/src/benchflow/sandbox/agentcore_provisioning.py index f8c518763..de9f25841 100644 --- a/src/benchflow/sandbox/agentcore_provisioning.py +++ b/src/benchflow/sandbox/agentcore_provisioning.py @@ -27,7 +27,8 @@ import hashlib import logging import re -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Iterator +from fnmatch import fnmatch from pathlib import Path from typing import Any @@ -81,23 +82,71 @@ def reset_cache() -> None: _KEY_LOCKS.clear() +def _dockerignore_matcher(context_dir: Path) -> Callable[[str], bool]: + """Compile ``.dockerignore`` into a predicate over context-relative paths. + + Docker semantics: one pattern per line, ``#`` comments, ``!`` re-includes, + and the last matching rule wins. ``fnmatch`` is close enough for the + patterns tasks actually use, with an added prefix rule so a directory + pattern excludes everything beneath it. + """ + ignore_file = context_dir / ".dockerignore" + if not ignore_file.is_file(): + return lambda _relative: False + + rules: list[tuple[str, bool]] = [] + for raw in ignore_file.read_text().splitlines(): + line = raw.strip() + if not line or line.startswith("#"): + continue + negated = line.startswith("!") + pattern = line.lstrip("!").strip().rstrip("/") + if pattern: + rules.append((pattern, negated)) + + def matches(relative: str) -> bool: + ignored = False + for pattern, negated in rules: + if fnmatch(relative, pattern) or relative.startswith(pattern + "/"): + ignored = not negated + return ignored + + return matches + + +def iter_context_files(context_dir: Path) -> Iterator[tuple[Path, str]]: + """The canonical Docker build context: ``(absolute path, relative path)``. + + Used by **both** the image digest and the CodeBuild upload so the two views + cannot drift. They previously did: the local Docker daemon honored + ``.dockerignore`` while the remote path zipped and uploaded every regular + file, which shipped ignored files — including secrets — into S3 and also + let an ignored file change the image identity. + + Symlinks are skipped so a task-controlled link cannot pull host files into + the image or the upload (#411). + """ + ignored = _dockerignore_matcher(context_dir) + for path in sorted(context_dir.rglob("*")): + if path.is_symlink() or not path.is_file(): + continue + relative = path.relative_to(context_dir).as_posix() + if relative in _GENERATED_NAMES or ignored(relative): + continue + yield path, relative + + def build_context_digest(context_dir: Path, dockerfile_text: str) -> str: """Content digest of everything that determines the built image. Hashes file *contents* rather than paths and mtimes on purpose: BenchFlow copies tasks into temporary directories before a run, so any identity based on location or timestamp would change every run and defeat image reuse - entirely. Symlinks are skipped for the same reason they are skipped on - upload (#411) — they must not pull host files into the image identity. + entirely. """ digest = hashlib.sha256() digest.update(dockerfile_text.encode()) - for path in sorted(context_dir.rglob("*")): - if path.is_symlink() or not path.is_file(): - continue - relative = path.relative_to(context_dir).as_posix() - if relative in _GENERATED_NAMES: - continue + for path, relative in iter_context_files(context_dir): digest.update(relative.encode()) digest.update(b"\0") digest.update(hashlib.sha256(path.read_bytes()).digest()) diff --git a/src/benchflow/sandbox/agentcore_reaper.py b/src/benchflow/sandbox/agentcore_reaper.py index 3abeeaf49..596874459 100644 --- a/src/benchflow/sandbox/agentcore_reaper.py +++ b/src/benchflow/sandbox/agentcore_reaper.py @@ -15,6 +15,7 @@ from __future__ import annotations import logging +from collections.abc import Mapping from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from typing import Any @@ -44,6 +45,25 @@ def summary(self) -> str: ) +def _runtime_timestamp(runtime: Mapping[str, Any]) -> datetime | None: + """Best available age signal for a runtime, or None if there is none. + + ``ListAgentRuntimes`` returns **only** ``lastUpdatedAt`` — there is no + ``createdAt`` in the list shape, though ``GetAgentRuntime`` has both. + Reading the wrong field silently yielded ``None`` for every runtime, which + skipped the age comparison entirely and made a one-day cleanup delete + minutes-old runtimes out from under a running matrix. + + Returning None here means "age unknown", and the caller must treat that as + not-stale rather than as stale. + """ + for field_name in ("createdAt", "lastUpdatedAt"): + value = runtime.get(field_name) + if isinstance(value, datetime): + return value if value.tzinfo else value.replace(tzinfo=UTC) + return None + + def _is_benchflow_managed(control: Any, arn: str) -> bool: """True only when the runtime carries BenchFlow's managed tag. @@ -80,12 +100,13 @@ def reap_stale_runtimes( report.skipped_unmanaged += 1 continue - created = runtime.get("createdAt") - if isinstance(created, datetime): - stamp = created if created.tzinfo else created.replace(tzinfo=UTC) - if stamp > cutoff: - report.skipped_recent += 1 - continue + stamp = _runtime_timestamp(runtime) + if stamp is None or stamp > cutoff: + # No usable age means we cannot prove the runtime is stale, and + # a runtime in use by a live matrix looks exactly like one that + # is idle. Fail closed. + report.skipped_recent += 1 + continue if dry_run: report.deleted.append(runtime_id) diff --git a/src/benchflow/sandbox/providers.py b/src/benchflow/sandbox/providers.py index 1ee2e34f8..f26fbcaa3 100644 --- a/src/benchflow/sandbox/providers.py +++ b/src/benchflow/sandbox/providers.py @@ -41,6 +41,9 @@ class SandboxProvider: #: registry instead of growing a ``sandbox == ""`` special case per #: backend that cannot isolate the network. enforces_no_network: bool = True + #: Whether the backend can run a task's docker-compose side services. + #: ``False`` means a multi-service task must be refused, not run partially. + supports_compose: bool = False @property def off_box_model(self) -> bool: @@ -51,11 +54,18 @@ def off_box_model(self) -> bool: # Ordered, docker-first. This tuple is the ONLY place the set is spelled out. _PROVIDERS: tuple[SandboxProvider, ...] = ( - SandboxProvider("docker", extra=None, model_proxy=ModelProxyLocation.HOST), + SandboxProvider( + "docker", + extra=None, + model_proxy=ModelProxyLocation.HOST, + supports_compose=True, + ), SandboxProvider( "daytona", extra="sandbox-daytona", model_proxy=ModelProxyLocation.SANDBOX, + # The DinD strategy runs compose inside the sandbox VM. + supports_compose=True, ), SandboxProvider( "modal", @@ -97,6 +107,10 @@ def off_box_model(self) -> bool: OFF_BOX_MODEL_PROVIDERS = SANDBOX_MODEL_PROXY_PROVIDERS +#: Providers that run exactly one container and cannot host compose services. +SINGLE_CONTAINER_PROVIDERS: frozenset[str] = frozenset( + p.name for p in _PROVIDERS if not p.supports_compose +) #: Providers that cannot enforce ``network_mode = "no-network"``. NO_NETWORK_UNSUPPORTED_PROVIDERS: frozenset[str] = frozenset( p.name for p in _PROVIDERS if not p.enforces_no_network diff --git a/src/benchflow/task/runtime_capabilities.py b/src/benchflow/task/runtime_capabilities.py index ba1f74d65..a77bec38b 100644 --- a/src/benchflow/task/runtime_capabilities.py +++ b/src/benchflow/task/runtime_capabilities.py @@ -15,9 +15,11 @@ from typing import cast from benchflow.rewards.rubric_config import criteria_aggregate_policy_from_rubric +from benchflow.sandbox._compose import compose_definition_path from benchflow.sandbox.providers import ( NO_NETWORK_UNSUPPORTED_PROVIDERS, SANDBOX_PROVIDER_SET, + SINGLE_CONTAINER_PROVIDERS, providers_phrase, ) from benchflow.task.config import ( @@ -514,6 +516,35 @@ def _append_layout_issues( sandbox=sandbox, ) _append_verifier_strategy_issue(unsupported, paths=paths, sandbox=sandbox) + _append_compose_issue(unsupported, task_dir=task_dir, sandbox=sandbox) + + +def _append_compose_issue( + unsupported: list[UnsupportedTaskFeature], + *, + task_dir: Path, + sandbox: str, +) -> None: + """Refuse a multi-service task on a backend that runs one container. + + Such a task would otherwise launch with its side services missing, and the + resulting failure would be attributed to the agent rather than reported as + an environment the backend cannot host. + """ + if sandbox not in SINGLE_CONTAINER_PROVIDERS: + return + compose = compose_definition_path(task_dir / "environment") + if compose is None: + return + _issue( + unsupported, + path=f"environment/{compose.name}", + reason=( + f"multi-service compose topologies are not supported by {sandbox}; " + "use the docker or daytona sandbox" + ), + sandbox=sandbox, + ) def _append_verifier_strategy_issue( diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py index d3ec90540..2902ac1da 100644 --- a/tests/test_agentcore_builder.py +++ b/tests/test_agentcore_builder.py @@ -173,3 +173,66 @@ def test_context_is_zipped_not_tarred(self, tmp_path): blob = builder._package(_request(tmp_path)) assert zipfile.is_zipfile(BytesIO(blob)) + + +class TestDockerIgnoreParity: + """The remote path must see the same context Docker would. + + The local daemon honors .dockerignore natively; the CodeBuild path builds + its own file list. When those diverged, ignored files — including secrets — + were zipped and uploaded to S3, and an ignored file also perturbed the + image identity. + """ + + def _context(self, tmp_path): + (tmp_path / "Dockerfile").write_text("FROM python:3.12-slim\n") + (tmp_path / "keep.txt").write_text("keep me") + (tmp_path / "secret.env").write_text("SUPER_SECRET=hunter2") + (tmp_path / "cache").mkdir() + (tmp_path / "cache" / "big.bin").write_text("x" * 64) + (tmp_path / ".dockerignore").write_text("secret.env\ncache/\n") + return tmp_path + + def test_ignored_files_are_not_uploaded(self, tmp_path): + context = self._context(tmp_path) + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + + blob = builder._package(_request(context)) + + with zipfile.ZipFile(BytesIO(blob)) as archive: + names = set(archive.namelist()) + payloads = b"".join(archive.read(n) for n in names) + + assert "secret.env" not in names + assert not any(n.startswith("cache/") for n in names) + assert b"hunter2" not in payloads + # Positive control: the archive is real and non-ignored files survive. + assert "keep.txt" in names + + def test_ignored_files_do_not_change_image_identity(self, tmp_path): + """An ignored file cannot affect the build, so it must not affect the tag.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + context = self._context(tmp_path) + before = provisioning.build_context_digest(context, "FROM python:3.12-slim\n") + + (context / "secret.env").write_text("SUPER_SECRET=rotated") + + assert ( + provisioning.build_context_digest(context, "FROM python:3.12-slim\n") + == before + ) + + def test_negation_re_includes_a_file(self, tmp_path): + """`!pattern` is Docker's re-include rule; last match wins.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + (tmp_path / "a.log").write_text("drop") + (tmp_path / "keep.log").write_text("keep") + (tmp_path / ".dockerignore").write_text("*.log\n!keep.log\n") + + relatives = {rel for _p, rel in provisioning.iter_context_files(tmp_path)} + + assert "a.log" not in relatives + assert "keep.log" in relatives diff --git a/tests/test_agentcore_parallel.py b/tests/test_agentcore_parallel.py index 26c6803e8..1baf98975 100644 --- a/tests/test_agentcore_parallel.py +++ b/tests/test_agentcore_parallel.py @@ -113,6 +113,11 @@ async def _build(self, image_uri, *, force_build): patch.object(AgentCoreSandbox, "_ensure_ecr_repository"), patch.object(AgentCoreSandbox, "_ecr_image_exists", return_value=False), patch.object(AgentCoreSandbox, "_build_and_push", new=_build), + patch.object( + AgentCoreSandbox, + "_resolve_image_digest", + lambda self, tag: f"reg/repo@sha256:{tag}", + ), ): sandboxes = [_sandbox(task_dir, session=f"run-{i}") for i in range(20)] uris = await asyncio.gather( @@ -157,6 +162,11 @@ async def _flaky(self, image_uri, *, force_build): patch.object(AgentCoreSandbox, "_ensure_ecr_repository"), patch.object(AgentCoreSandbox, "_ecr_image_exists", return_value=False), patch.object(AgentCoreSandbox, "_build_and_push", new=_flaky), + patch.object( + AgentCoreSandbox, + "_resolve_image_digest", + lambda self, tag: f"reg/repo@sha256:{tag}", + ), ): env = _sandbox(task_dir) with pytest.raises(RuntimeError, match="transient throttle"): @@ -354,3 +364,211 @@ def test_daytona_failure_does_not_block_agentcore_cleanup(self, monkeypatch): sandbox_cli.sandbox_cleanup(dry_run=True, max_age_minutes=60) assert calls == ["agentcore"] + + +class TestAwsResponseShapeConformance: + """Pin the mocked shapes to the SDK's real ones. + + Every fixture in this file invents AWS responses. When an invented shape + drifts from the service model the tests keep passing while production + breaks — which is exactly how the reaper came to read a ``createdAt`` that + ``ListAgentRuntimes`` never returns, and so deleted fresh runtimes. + """ + + def _shape(self, operation, member=None): + boto3 = pytest.importorskip("boto3") + client = boto3.Session(region_name="us-west-2").client( + "bedrock-agentcore-control", + aws_access_key_id="x", + aws_secret_access_key="y", + ) + shape = client.meta.service_model.operation_model(operation).output_shape + return shape.members[member].member if member else shape + + def test_list_agent_runtimes_has_no_created_at(self): + """The field the reaper originally keyed on does not exist here.""" + members = self._shape("ListAgentRuntimes", "agentRuntimes").members + + assert "lastUpdatedAt" in members + assert "createdAt" not in members + + def test_get_agent_runtime_carries_the_bound_image(self): + """The adoption check reads this to refuse a mismatched runtime.""" + members = self._shape("GetAgentRuntime").members + + assert "agentRuntimeArtifact" in members + container = members["agentRuntimeArtifact"].members["containerConfiguration"] + assert "containerUri" in container.members + + +class TestReaperAgeHandling: + def _control(self, runtimes, tags): + control = MagicMock() + paginator = MagicMock() + paginator.paginate.return_value = [{"agentRuntimes": runtimes}] + control.get_paginator.return_value = paginator + control.list_tags_for_resource.side_effect = lambda resourceArn: { + "tags": tags.get(resourceArn, {}) + } + return control + + def _managed(self, arn): + return {arn: {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}} + + def test_fresh_runtime_in_the_real_list_shape_is_kept(self): + """Reproduces the reported P1 with the shape AWS actually returns. + + ``ListAgentRuntimes`` has only ``lastUpdatedAt``; keying on + ``createdAt`` yielded None, skipped the age comparison, and selected a + minutes-old runtime for deletion under a one-day policy. + """ + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + runtimes = [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "agentRuntimeName": "bf_x", + "lastUpdatedAt": datetime.now(UTC) - timedelta(minutes=3), + "status": "READY", + } + ] + control = self._control(runtimes, self._managed("arn:mine")) + + report = reap_stale_runtimes(control, max_age_minutes=1440, dry_run=True) + + assert report.deleted == [] + assert report.skipped_recent == 1 + + def test_runtime_with_no_timestamp_is_kept(self): + """Unknown age must never be read as 'old enough to delete'.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control( + [{"agentRuntimeArn": "arn:mine", "agentRuntimeId": "mine"}], + self._managed("arn:mine"), + ) + + report = reap_stale_runtimes(control, max_age_minutes=1440) + + assert report.deleted == [] + assert report.skipped_recent == 1 + control.delete_agent_runtime.assert_not_called() + + def test_genuinely_old_runtime_in_the_real_shape_is_reaped(self): + """The positive control: cleanup must still do its job.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control( + [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "lastUpdatedAt": datetime.now(UTC) - timedelta(days=3), + } + ], + self._managed("arn:mine"), + ) + + report = reap_stale_runtimes(control, max_age_minutes=1440) + + assert report.deleted == ["mine"] + + +class TestRuntimeImageBinding: + @pytest.mark.asyncio + async def test_a_runtime_bound_to_another_image_is_updated( + self, tmp_path, monkeypatch + ): + """Adopting a stale runtime would run the agent in the wrong environment.""" + from botocore.exceptions import ClientError + + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + env = _sandbox(_make_task(tmp_path)) + control = MagicMock() + control.create_agent_runtime.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "exists"}}, + "CreateAgentRuntime", + ) + control.get_agent_runtime.return_value = { + "status": "READY", + "agentRuntimeArtifact": { + "containerConfiguration": {"containerUri": "new.example/repo@sha256:b"} + }, + } + + with ( + patch.object(env, "_client", return_value=control), + patch.object( + provisioning, + "find_runtime_by_name", + return_value=("arn:rt", "rt-1", "old.example/repo@sha256:a"), + ), + ): + arn, _rid = await env._create_or_adopt_runtime( + "bf_x", "new.example/repo@sha256:b" + ) + + assert arn == "arn:rt" + control.update_agent_runtime.assert_called_once() + sent = control.update_agent_runtime.call_args.kwargs + assert ( + sent["agentRuntimeArtifact"]["containerConfiguration"]["containerUri"] + == "new.example/repo@sha256:b" + ) + + def test_adoption_fails_closed_when_the_image_still_mismatches(self): + """If the update did not take, refuse rather than run the wrong image.""" + from benchflow.sandbox.agentcore import AgentCoreSandbox + from benchflow.sandbox.protocol import SandboxStartupError + + control = MagicMock() + control.get_agent_runtime.return_value = { + "agentRuntimeArtifact": { + "containerConfiguration": {"containerUri": "old.example/repo@sha256:a"} + } + } + + with pytest.raises(SandboxStartupError, match="wrong image"): + AgentCoreSandbox._verify_runtime_image( + control, "rt-1", "new.example/repo@sha256:b" + ) + + +class TestComposeRejection: + def test_compose_task_is_refused_at_construction(self, tmp_path): + """One container cannot host a task's side services.""" + task_dir = _make_task(tmp_path) + (task_dir / "docker-compose.yaml").write_text("services:\n target: {}\n") + + with pytest.raises(ValueError, match="single-container"): + _sandbox(task_dir) + + def test_compose_task_is_refused_by_the_capability_gate(self, tmp_path): + """Fail during planning, before any image is built.""" + from benchflow.task.config import TaskConfig + from benchflow.task.runtime_capabilities import validate_task_runtime_support + + (tmp_path / "environment").mkdir() + (tmp_path / "environment" / "Dockerfile").write_text("FROM scratch\n") + (tmp_path / "environment" / "compose.yaml").write_text("services: {}\n") + + issues = validate_task_runtime_support( + TaskConfig.model_validate({}), sandbox="agentcore", task_dir=tmp_path + ) + + assert any("compose" in issue.reason for issue in issues) + + def test_docker_still_accepts_compose_tasks(self, tmp_path): + """The gate must not regress the backends that do support compose.""" + from benchflow.task.config import TaskConfig + from benchflow.task.runtime_capabilities import validate_task_runtime_support + + (tmp_path / "environment").mkdir() + (tmp_path / "environment" / "docker-compose.yaml").write_text("services: {}\n") + + issues = validate_task_runtime_support( + TaskConfig.model_validate({}), sandbox="docker", task_dir=tmp_path + ) + + assert not any("compose" in issue.reason for issue in issues) diff --git a/tests/test_agentcore_sandbox.py b/tests/test_agentcore_sandbox.py index 2c46a342d..ead33902e 100644 --- a/tests/test_agentcore_sandbox.py +++ b/tests/test_agentcore_sandbox.py @@ -21,7 +21,10 @@ import base64 import contextlib import os +import re +import shlex import tarfile +from io import BytesIO from unittest.mock import MagicMock, patch import pytest @@ -254,6 +257,24 @@ async def _record(command, **kwargs): assert len(commands) < 13 assert any("tar -xzf" in c for c in commands) + @staticmethod + def _extract_staged_archive(commands: list[str]) -> tarfile.TarFile: + """Rebuild the tar that upload_dir actually streamed to the sandbox. + + ``_upload_via_tar`` base64-encodes the archive into ``printf`` commands, + so asserting against the raw command text can never observe the + archive's contents — a regression that started following symlinks would + have passed. Decode the staged chunks back into the real tar instead. + """ + encoded = "".join( + match.group(1) + for command in commands + if (match := re.search(r"printf %s (\S+) >>? /tmp/", command)) + ) + # shlex.quote may wrap the chunk; strip the quoting before decoding. + encoded = "".join(shlex.split(encoded)) if encoded else "" + return tarfile.open(fileobj=BytesIO(base64.b64decode(encoded)), mode="r:gz") + @pytest.mark.asyncio async def test_upload_dir_skips_symlinks(self, sandbox, tmp_path): """Guards #411: a task symlink must not exfiltrate host files.""" @@ -273,8 +294,16 @@ async def _record(command, **kwargs): with patch.object(sandbox, "exec", side_effect=_record): await sandbox.upload_dir(source, "/workspace") - blob = "".join(staged) - assert "do not ship me" not in blob + with self._extract_staged_archive(staged) as tar: + names = tar.getnames() + payloads = { + name: (tar.extractfile(name) or BytesIO()).read() for name in names + } + + # Positive control: the archive really was inspected, not empty. + assert any(name.endswith("real.txt") for name in names) + assert not any(name.endswith("link.txt") for name in names) + assert all(b"do not ship me" not in blob for blob in payloads.values()) @pytest.mark.asyncio async def test_download_dir_extracts_the_returned_archive(self, sandbox, tmp_path): From 39e2f09f7b3b42b9bda9bc695cc80aae68bd425d Mon Sep 17 00:00:00 2001 From: Bingran You Date: Mon, 27 Jul 2026 20:10:50 -0700 Subject: [PATCH 05/13] fix(agentcore): lease-gated cleanup, Docker-compatible ignore, transport wake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second review round: seven P1 and four P2 findings, each re-verified live. P1 — cleanup could delete a runtime serving an active session. Session traffic does not move a runtime's `lastUpdatedAt`, and there is no API that enumerates a runtime's active sessions (`ListSessions` is Memory-scoped), so control-plane age cannot distinguish idle from in-use. Added an explicit lease: provisioning tags the runtime with an expiry covering the longest a session started then could live, and cleanup refuses any runtime whose lease is unexpired, unparseable, or unreadable. Written once per runtime per process, so it costs nothing per rollout. Live: age-0 cleanup now reports active=1, deleted=0 for a runtime with a live session. P1 — negative cleanup ages selected fresh runtimes. A negative age puts the cutoff in the future, so everything looks stale. Rejected in the library entry point and constrained on the Typer option. P1 — `.dockerignore` matching leaked root secrets to S3. `fnmatch` does not implement Docker's semantics: `/secret.env` never matched because the leading separator was kept, and `**/*.pem` could not span zero directories. Replaced with a real pattern translator (separators stripped, `*`/`?` bounded by `/`, `**` spanning zero or more segments, ancestor-directory matches excluding contents). Live: both patterns now exclude their files. P1 — the runtime single-flight key ignored the image. A force rebuild pushes new bytes under the same context identity, so the memo returned the old ARN and bypassed compare/update/verify entirely. The key now includes the immutable image URI. P1 — adoption reset lifecycle to service defaults. `update_agent_runtime` omitted `lifecycleConfiguration`, so a configured 600/7200 window came back as 900/28800 and would reclaim long sessions early. Now passed on update and verified alongside the image; drift fails closed. Live: 600/7200 preserved. P1 — image identity omitted the shim and mode bits. The generated shim is copied into the image, and an entrypoint flipped 0644 to 0755 changes the container, yet neither moved the digest — so an upgrade could reuse a stale image. Both are hashed now. P1 — a dead shell reader hung until the 900s read timeout while reporting itself running. The reader now enqueues an end sentinel in a `finally`, so readline reports the disconnect immediately, and liveness includes reader state. Clean EOF is fatal too, not just an exception. P2 — the remote size gate compared floored megabytes, so cap-plus-one-byte passed the gate that `image_size_error()` rejects; it compares raw bytes now. P2 — every shell channel fed the ACP stream, so a STDERR diagnostic could be parsed as JSON-RPC; side channels are now logged separately, with unknown channels still treated as agent output so an SDK mismatch cannot mute the agent. P2 — bucket hardening was skipped for existing buckets and swallowed every error; access controls and expiry are now asserted on every call and failures are fatal. P2 — `_ecr_image_exists` turned access-denied and throttling into cache misses, starting doomed builds; only a genuine not-found is a miss. Known limit: the lease covers the configured session window from provisioning time. A run longer than BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC needs that raised, or cleanup deferred. --- src/benchflow/cli/sandbox.py | 7 +- src/benchflow/sandbox/agentcore.py | 92 ++++++++- src/benchflow/sandbox/agentcore_builder.py | 108 +++++----- .../sandbox/agentcore_provisioning.py | 70 +++++-- src/benchflow/sandbox/agentcore_reaper.py | 64 +++++- src/benchflow/sandbox/process/agentcore.py | 82 ++++++-- tests/test_agentcore_builder.py | 18 +- tests/test_agentcore_parallel.py | 194 +++++++++++++++++- tests/test_agentcore_transport.py | 143 +++++++++++++ 9 files changed, 675 insertions(+), 103 deletions(-) create mode 100644 tests/test_agentcore_transport.py diff --git a/src/benchflow/cli/sandbox.py b/src/benchflow/cli/sandbox.py index 53496412c..570cb4353 100644 --- a/src/benchflow/cli/sandbox.py +++ b/src/benchflow/cli/sandbox.py @@ -231,7 +231,12 @@ def sandbox_cleanup_cmd( bool, typer.Option("--dry-run", help="List sandboxes without deleting") ] = False, max_age_minutes: Annotated[ - int, typer.Option("--max-age", help="Delete sandboxes older than N minutes") + int, + typer.Option( + "--max-age", + min=0, + help="Delete sandboxes older than N minutes", + ), ] = 1440, ) -> None: """Clean up orphaned Daytona sandboxes.""" diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py index e70f24516..2507d52db 100644 --- a/src/benchflow/sandbox/agentcore.py +++ b/src/benchflow/sandbox/agentcore.py @@ -51,6 +51,7 @@ import tarfile import threading import uuid +from datetime import UTC, datetime, timedelta from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING, Any @@ -296,7 +297,7 @@ def _image_identity(self) -> tuple[str, str]: running a build to discover it. """ digest = provisioning.build_context_digest( - self.environment_dir, self._generated_dockerfile_text() + self.environment_dir, self._generated_dockerfile_text(), _PING_SHIM ) return digest, provisioning.image_tag(self.environment_name, digest) @@ -396,6 +397,11 @@ def _ensure_ecr_repository(self) -> None: raise def _ecr_image_exists(self, tag: str) -> bool: + """True when *tag* is present. Only a real miss counts as a miss. + + Treating access-denied or throttling as "not found" would start a + doomed build and bury the actual AWS error. + """ from botocore.exceptions import ClientError try: @@ -403,8 +409,13 @@ def _ecr_image_exists(self, tag: str) -> bool: repositoryName=self._ecr_repository, imageIds=[{"imageTag": tag}] ) return True - except ClientError: - return False + except ClientError as exc: + if exc.response["Error"]["Code"] in { + "ImageNotFoundException", + "RepositoryNotFoundException", + }: + return False + raise async def _ensure_runtime(self, image_uri: str) -> str: """Resolve the shared agent runtime for *image_uri*, creating it once. @@ -422,7 +433,12 @@ async def _ensure_runtime(self, image_uri: str) -> str: async def _create() -> tuple[str, str]: return await self._create_or_adopt_runtime(name, image_uri) - arn, runtime_id = await provisioning.once(f"runtime:{name}", _create) + # Keyed on the image as well as the name: with only the name, a + # force-rebuild that pushes new bytes under the same context identity + # hits the memo and silently skips compare/update/verify. + arn, runtime_id = await provisioning.once( + f"runtime:{name}:{image_uri}", _create + ) self._runtime_id = runtime_id return arn @@ -465,6 +481,7 @@ async def _create_or_adopt_runtime( if existing is None: raise arn, runtime_id, current_image = existing + await asyncio.to_thread(self._write_lease, control, arn) if current_image != image_uri: # Adopting a runtime bound to different bytes would silently run # the wrong environment. The name encodes the build-context @@ -486,15 +503,24 @@ async def _create_or_adopt_runtime( roleArn=self._require_role_arn(), networkConfiguration={"networkMode": "PUBLIC"}, protocolConfiguration=_PROTOCOL_CONFIGURATION, + # Omitting this resets idle/max-lifetime to the service + # defaults, silently discarding the caller's configured + # window and reclaiming long sessions early. + lifecycleConfiguration=self._lifecycle_configuration(), ) await asyncio.to_thread(self._wait_ready, control, runtime_id) await asyncio.to_thread( - self._verify_runtime_image, control, runtime_id, image_uri + self._verify_adopted_runtime, + control, + runtime_id, + image_uri, + self._lifecycle_configuration(), ) self.logger.info("Adopted existing AgentCore runtime %s", arn) return arn, runtime_id runtime_id = created["agentRuntimeId"] + await asyncio.to_thread(self._write_lease, control, created["agentRuntimeArn"]) await asyncio.to_thread(self._wait_ready, control, runtime_id) self.logger.info( "Registered AgentCore runtime %s for %s", @@ -503,13 +529,46 @@ async def _create_or_adopt_runtime( ) return created["agentRuntimeArn"], runtime_id + def _write_lease(self, control: Any, runtime_arn: str) -> None: + """Mark the runtime as possibly-in-use until the session window closes. + + There is no API that enumerates a runtime's active sessions + (``ListSessions`` is Memory-scoped), and session traffic does not move + the runtime's ``lastUpdatedAt`` — so control-plane age alone cannot + distinguish an idle runtime from one serving a matrix right now. + + The lease is the explicit contract that closes that gap: it extends to + the longest a session started now could still be alive, and cleanup + refuses to touch a runtime whose lease has not expired. Written once + per runtime per process (provisioning is single-flighted), so it costs + nothing per rollout. + """ + lifecycle = self._lifecycle_configuration() + window = max(lifecycle["maxLifetime"], lifecycle["idleRuntimeSessionTimeout"]) + until = datetime.now(UTC) + timedelta(seconds=window) + try: + control.tag_resource( + resourceArn=runtime_arn, + tags={provisioning.LEASE_TAG: until.isoformat()}, + ) + except Exception as exc: + # A runtime without a readable lease is treated as in-use by the + # reaper, so a failure here is conservative, not dangerous. + self.logger.warning("Could not write AgentCore lease: %s", exc) + @staticmethod - def _verify_runtime_image(control: Any, runtime_id: str, image_uri: str) -> None: - """Fail closed if the runtime is not bound to the image we expect. + def _verify_adopted_runtime( + control: Any, + runtime_id: str, + image_uri: str, + lifecycle: dict[str, int], + ) -> None: + """Fail closed unless the adopted runtime matches what this run needs. - The adoption path is the one place a runtime can predate this process, - so it is the one place a mismatch could otherwise go unnoticed and run - an agent against the wrong environment. + Adoption is the one path where a runtime can predate this process, so + it is the one path where a mismatch could otherwise go unnoticed — an + agent run against the wrong image, or against a shorter session window + than the caller configured. """ detail = control.get_agent_runtime(agentRuntimeId=runtime_id) artifact = detail.get("agentRuntimeArtifact") or {} @@ -521,6 +580,19 @@ def _verify_runtime_image(control: Any, runtime_id: str, image_uri: str) -> None "image.", sandbox_id=runtime_id, ) + actual = detail.get("lifecycleConfiguration") or {} + drifted = { + key: (actual.get(key), value) + for key, value in lifecycle.items() + if actual.get(key) != value + } + if drifted: + raise SandboxStartupError( + f"AgentCore runtime {runtime_id} lifecycle is {actual}, not the " + f"configured {lifecycle} (drift: {drifted}). Sessions could be " + "reclaimed earlier than this run expects.", + sandbox_id=runtime_id, + ) def _lifecycle_configuration(self) -> dict[str, int]: """Session idle/lifetime caps. diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py index f05cc0fb6..f51f1f0ba 100644 --- a/src/benchflow/sandbox/agentcore_builder.py +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -211,10 +211,13 @@ def _login(self, registry: str) -> None: f"docker build --platform linux/arm64 " f"-f {provisioning.GENERATED_DOCKERFILE} -t $BF_IMAGE_URI .", 'SIZE=$(docker image inspect -f "{{.Size}}" $BF_IMAGE_URI)', - "MB=$((SIZE / 1048576))", - 'echo "benchflow: image size ${MB} MB"', - f'if [ "$MB" -gt {provisioning.MAX_IMAGE_MB} ]; then ' - f'echo "BENCHFLOW_IMAGE_TOO_LARGE ${{MB}}"; exit 1; fi', + 'echo "benchflow: image size $((SIZE / 1048576)) MB"', + # Compare raw bytes, not floored megabytes: an image of + # exactly the cap plus one byte floors to the cap and would + # slip past a megabyte comparison that image_size_error() + # correctly rejects. + f'if [ "$SIZE" -gt {provisioning.MAX_IMAGE_MB * 1024 * 1024} ]; ' + 'then echo "BENCHFLOW_IMAGE_TOO_LARGE $SIZE"; exit 1; fi', "docker push $BF_IMAGE_URI", ] }, @@ -298,61 +301,64 @@ def _package(self, request: BuildRequest) -> bytes: return buffer.getvalue() def _ensure_bucket(self) -> None: + """Create the build bucket if needed, and always assert its controls. + + Hardening runs on every call, not only on creation. An existing bucket + — one a previous run created before these controls existed, or one the + operator pointed us at — would otherwise keep uploaded build contexts + public-by-default and un-expiring forever. Build contexts can contain + task source and credentials, so failing to establish those controls is + a stop condition, not a warning. + """ from botocore.exceptions import ClientError s3 = self._client("s3") try: s3.head_bucket(Bucket=self._bucket) - return - except ClientError: - pass - kwargs: dict[str, Any] = {"Bucket": self._bucket} - if self._region != "us-east-1": - kwargs["CreateBucketConfiguration"] = {"LocationConstraint": self._region} - try: - s3.create_bucket(**kwargs) except ClientError as exc: - if exc.response["Error"]["Code"] not in { - "BucketAlreadyOwnedByYou", - "BucketAlreadyExists", - }: + code = str(exc.response["Error"].get("Code", "")) + if code not in {"404", "NoSuchBucket", "NotFound"}: + # 403 and throttling are not "absent" — creating over them + # would fail confusingly and mask the real problem. raise - for call, kw in ( - ( - s3.put_public_access_block, - { - "Bucket": self._bucket, - "PublicAccessBlockConfiguration": { - "BlockPublicAcls": True, - "IgnorePublicAcls": True, - "BlockPublicPolicy": True, - "RestrictPublicBuckets": True, - }, - }, - ), - ( - s3.put_bucket_lifecycle_configuration, - { - "Bucket": self._bucket, - # Build contexts are consumed within minutes; expiring them - # keeps a benchmark run from accruing storage cost forever. - "LifecycleConfiguration": { - "Rules": [ - { - "ID": "benchflow-expire-build-contexts", - "Status": "Enabled", - "Filter": {"Prefix": "contexts/"}, - "Expiration": {"Days": 1}, - } - ] - }, - }, - ), - ): + kwargs: dict[str, Any] = {"Bucket": self._bucket} + if self._region != "us-east-1": + kwargs["CreateBucketConfiguration"] = { + "LocationConstraint": self._region + } try: - call(**kw) - except Exception as exc: - logger.debug("Bucket hardening step skipped: %s", exc) + s3.create_bucket(**kwargs) + except ClientError as create_exc: + if create_exc.response["Error"]["Code"] not in { + "BucketAlreadyOwnedByYou", + "BucketAlreadyExists", + }: + raise + + s3.put_public_access_block( + Bucket=self._bucket, + PublicAccessBlockConfiguration={ + "BlockPublicAcls": True, + "IgnorePublicAcls": True, + "BlockPublicPolicy": True, + "RestrictPublicBuckets": True, + }, + ) + s3.put_bucket_lifecycle_configuration( + Bucket=self._bucket, + # Build contexts are consumed within minutes; expiring them keeps a + # benchmark run from accruing storage cost — and exposure — forever. + LifecycleConfiguration={ + "Rules": [ + { + "ID": "benchflow-expire-build-contexts", + "Status": "Enabled", + "Filter": {"Prefix": "contexts/"}, + "Expiration": {"Days": 1}, + } + ] + }, + ) def _ensure_project(self) -> None: from botocore.exceptions import ClientError diff --git a/src/benchflow/sandbox/agentcore_provisioning.py b/src/benchflow/sandbox/agentcore_provisioning.py index de9f25841..5420fc45e 100644 --- a/src/benchflow/sandbox/agentcore_provisioning.py +++ b/src/benchflow/sandbox/agentcore_provisioning.py @@ -28,7 +28,6 @@ import logging import re from collections.abc import Awaitable, Callable, Iterator -from fnmatch import fnmatch from pathlib import Path from typing import Any @@ -44,6 +43,12 @@ # resources apart from anything else in the account. MANAGED_TAG = "benchflow-managed" MANAGED_VALUE = "1" +#: Tag holding an ISO-8601 instant until which a runtime must not be reaped. +#: There is no API to enumerate a runtime's active sessions — ``ListSessions`` +#: is Memory-scoped — and session traffic does not move the runtime's +#: ``lastUpdatedAt``. A lease written at provisioning time is therefore the +#: only signal cleanup has that a runtime may still be serving a matrix. +LEASE_TAG = "benchflow-lease-until" # Not adjustable per AWS service quotas: "Maximum size (in MB) for a Docker # image in an AgentCore Runtime" = 2048. @@ -82,32 +87,63 @@ def reset_cache() -> None: _KEY_LOCKS.clear() +def _translate_ignore_pattern(pattern: str) -> re.Pattern[str]: + """Compile one ``.dockerignore`` pattern to a regex over relative paths. + + Follows Docker's matcher rather than approximating it with ``fnmatch``: + leading and trailing separators are stripped (so ``/secret.env`` is the + root file, not an absolute path that never matches), ``*`` and ``?`` stop + at a separator, and ``**`` spans zero or more path segments. Approximating + these let ``/secret.env`` and ``**/*.pem`` through, which uploaded root + secrets to S3. + """ + cleaned = pattern.strip().strip("/") + segments = cleaned.split("/") + regex = "" + for index, segment in enumerate(segments): + last = index == len(segments) - 1 + if segment == "**": + regex += "(?:.*)?" if last else "(?:[^/]+/)*" + continue + for char in segment: + if char == "*": + regex += "[^/]*" + elif char == "?": + regex += "[^/]" + else: + regex += re.escape(char) + if not last: + regex += "/" + return re.compile(f"^{regex}$") + + def _dockerignore_matcher(context_dir: Path) -> Callable[[str], bool]: """Compile ``.dockerignore`` into a predicate over context-relative paths. - Docker semantics: one pattern per line, ``#`` comments, ``!`` re-includes, - and the last matching rule wins. ``fnmatch`` is close enough for the - patterns tasks actually use, with an added prefix rule so a directory - pattern excludes everything beneath it. + Docker semantics: ``#`` comments, ``!`` re-includes, last matching rule + wins, and a rule matching a directory excludes everything beneath it. """ ignore_file = context_dir / ".dockerignore" if not ignore_file.is_file(): return lambda _relative: False - rules: list[tuple[str, bool]] = [] + rules: list[tuple[re.Pattern[str], bool]] = [] for raw in ignore_file.read_text().splitlines(): line = raw.strip() if not line or line.startswith("#"): continue negated = line.startswith("!") - pattern = line.lstrip("!").strip().rstrip("/") - if pattern: - rules.append((pattern, negated)) + body = line.lstrip("!").strip() + if body.strip("/"): + rules.append((_translate_ignore_pattern(body), negated)) def matches(relative: str) -> bool: + # A rule that matches an ancestor directory excludes its contents. + segments = relative.split("/") + candidates = ["/".join(segments[: i + 1]) for i in range(len(segments))] ignored = False - for pattern, negated in rules: - if fnmatch(relative, pattern) or relative.startswith(pattern + "/"): + for compiled, negated in rules: + if any(compiled.match(candidate) for candidate in candidates): ignored = not negated return ignored @@ -136,19 +172,29 @@ def iter_context_files(context_dir: Path) -> Iterator[tuple[Path, str]]: yield path, relative -def build_context_digest(context_dir: Path, dockerfile_text: str) -> str: +def build_context_digest( + context_dir: Path, dockerfile_text: str, shim_text: str = "" +) -> str: """Content digest of everything that determines the built image. Hashes file *contents* rather than paths and mtimes on purpose: BenchFlow copies tasks into temporary directories before a run, so any identity based on location or timestamp would change every run and defeat image reuse entirely. + + ``shim_text`` and the executable mode bit are part of the identity because + both change the built image: the shim is copied in as the entrypoint, and + an ``entrypoint.sh`` flipped from 0644 to 0755 produces a different + container even though every byte of content is unchanged. """ digest = hashlib.sha256() digest.update(dockerfile_text.encode()) + digest.update(b"\0shim\0") + digest.update(shim_text.encode()) for path, relative in iter_context_files(context_dir): digest.update(relative.encode()) digest.update(b"\0") + digest.update(b"x" if path.stat().st_mode & 0o111 else b"-") digest.update(hashlib.sha256(path.read_bytes()).digest()) return digest.hexdigest() diff --git a/src/benchflow/sandbox/agentcore_reaper.py b/src/benchflow/sandbox/agentcore_reaper.py index 596874459..f2ec670af 100644 --- a/src/benchflow/sandbox/agentcore_reaper.py +++ b/src/benchflow/sandbox/agentcore_reaper.py @@ -20,7 +20,11 @@ from datetime import UTC, datetime, timedelta from typing import Any -from benchflow.sandbox.agentcore_provisioning import MANAGED_TAG, MANAGED_VALUE +from benchflow.sandbox.agentcore_provisioning import ( + LEASE_TAG, + MANAGED_TAG, + MANAGED_VALUE, +) logger = logging.getLogger("benchflow").getChild("agentcore-reaper") @@ -35,13 +39,14 @@ class ReapReport: deleted: list[str] = field(default_factory=list) skipped_unmanaged: int = 0 skipped_recent: int = 0 + skipped_active: int = 0 errors: list[str] = field(default_factory=list) def summary(self) -> str: return ( f"scanned={self.scanned} deleted={len(self.deleted)} " f"unmanaged={self.skipped_unmanaged} recent={self.skipped_recent} " - f"errors={len(self.errors)}" + f"active={self.skipped_active} errors={len(self.errors)}" ) @@ -64,18 +69,37 @@ def _runtime_timestamp(runtime: Mapping[str, Any]) -> datetime | None: return None -def _is_benchflow_managed(control: Any, arn: str) -> bool: - """True only when the runtime carries BenchFlow's managed tag. +def _read_tags(control: Any, arn: str) -> dict[str, str] | None: + """Tags for *arn*, or None when they cannot be read. - Fails closed: if the tags cannot be read, the runtime is treated as - someone else's and left alone. + None means "cannot prove anything about this runtime", which every caller + must treat as a reason to leave it alone. """ try: - tags = control.list_tags_for_resource(resourceArn=arn).get("tags", {}) + return control.list_tags_for_resource(resourceArn=arn).get("tags", {}) except Exception: logger.debug("Could not read tags for %s; leaving it alone", arn) + return None + + +def _lease_is_active(tags: Mapping[str, str], now: datetime) -> bool: + """Whether a runtime is still leased, and so must not be deleted. + + A malformed lease counts as active. Cleanup is destructive and there is no + API to enumerate active sessions, so an unparseable lease is a reason to + stop, not to proceed. + """ + raw = tags.get(LEASE_TAG) + if not raw: return False - return tags.get(MANAGED_TAG) == MANAGED_VALUE + try: + until = datetime.fromisoformat(raw) + except ValueError: + logger.warning("Unparseable AgentCore lease %r; treating as active", raw) + return True + if until.tzinfo is None: + until = until.replace(tzinfo=UTC) + return until > now def reap_stale_runtimes( @@ -85,9 +109,23 @@ def reap_stale_runtimes( dry_run: bool = False, now: datetime | None = None, ) -> ReapReport: - """Delete BenchFlow-managed runtimes older than *max_age_minutes*.""" + """Delete BenchFlow-managed runtimes older than *max_age_minutes*. + + Two independent gates must both pass: the runtime's lease must have + expired, and its control-plane age must exceed *max_age_minutes*. The lease + exists because session traffic does not move ``lastUpdatedAt``, so age + alone cannot tell an idle runtime from one serving a matrix right now. + """ + if max_age_minutes < 0: + # A negative age puts the cutoff in the future, which makes every + # runtime — including one created a second ago — look stale. + raise ValueError( + f"max_age_minutes must be >= 0, got {max_age_minutes}. A negative " + "age would select every runtime, including active ones." + ) report = ReapReport() - cutoff = (now or datetime.now(UTC)) - timedelta(minutes=max_age_minutes) + moment = now or datetime.now(UTC) + cutoff = moment - timedelta(minutes=max_age_minutes) for page in control.get_paginator("list_agent_runtimes").paginate(): for runtime in page.get("agentRuntimes", []): @@ -96,9 +134,13 @@ def reap_stale_runtimes( runtime_id = runtime.get("agentRuntimeId") if not arn or not runtime_id: continue - if not _is_benchflow_managed(control, arn): + tags = _read_tags(control, arn) + if tags is None or tags.get(MANAGED_TAG) != MANAGED_VALUE: report.skipped_unmanaged += 1 continue + if _lease_is_active(tags, moment): + report.skipped_active += 1 + continue stamp = _runtime_timestamp(runtime) if stamp is None or stamp > cutoff: diff --git a/src/benchflow/sandbox/process/agentcore.py b/src/benchflow/sandbox/process/agentcore.py index 2e78cdcfb..deee6a00e 100644 --- a/src/benchflow/sandbox/process/agentcore.py +++ b/src/benchflow/sandbox/process/agentcore.py @@ -27,6 +27,13 @@ logger = logging.getLogger(__name__) +#: Channels the AgentCore shell multiplexes alongside agent stdout. Forwarding +#: these into the ACP stream would let a diagnostic be parsed as JSON-RPC. +_SIDE_CHANNELS = ("STDERR", "STATUS", "CLOSE") + +#: Queue sentinel meaning "the frame reader has ended"; see _drain_frames. +_READER_ENDED = object() + _START_MARKER_TIMEOUT_SEC = 180 _READLINE_TIMEOUT_ENV = "BENCHFLOW_AGENTCORE_READLINE_TIMEOUT" _READLINE_TIMEOUT_DEFAULT_SEC = 900.0 @@ -67,10 +74,11 @@ def __init__( self._region = region self._shell: Any = None self._reader_task: asyncio.Task[None] | None = None - self._line_buffer: asyncio.Queue[bytes] = asyncio.Queue() + self._line_buffer: asyncio.Queue[Any] = asyncio.Queue() self._partial = b"" self._closed = False self._failure: BaseException | None = None + self._reader_done = False @classmethod def from_sandbox_env(cls, env: Any) -> AgentCoreProcess: @@ -87,9 +95,18 @@ def from_sandbox_env(cls, env: Any) -> AgentCoreProcess: ) async def _drain_frames(self) -> None: - """Frame WebSocket payloads into newline-terminated lines.""" + """Frame STDOUT payloads into newline-terminated lines. + + Only the STDOUT channel becomes ACP input. The shell multiplexes + STDERR and lifecycle channels over the same socket, and forwarding + those verbatim would let a diagnostic line enter the JSON-RPC stream + and corrupt — or impersonate — protocol traffic. + """ try: async for frame in self._shell: + if not self._is_stdout(frame): + self._log_non_stdout(frame) + continue payload = frame.payload if isinstance(payload, str): payload = payload.encode() @@ -102,11 +119,40 @@ async def _drain_frames(self) -> None: except asyncio.CancelledError: raise except BaseException as exc: - # Record rather than raise: this runs in a background task whose - # exception would otherwise be swallowed, leaving readline() to - # hang until its timeout instead of reporting the real cause. self._failure = exc logger.warning("AgentCore shell reader stopped: %s", exc) + finally: + # Wake any waiting readline() right now. Without this sentinel a + # dead transport surfaces only when the read timeout expires, so an + # infrastructure disconnect becomes a 15-minute silent hang. + self._reader_done = True + with contextlib.suppress(Exception): + self._line_buffer.put_nowait(_READER_ENDED) + + @staticmethod + def _is_stdout(frame: Any) -> bool: + """Whether *frame* carries agent stdout rather than a side channel. + + Side channels are named explicitly and everything else is treated as + agent output. The inverse — allowing only a recognized ``STDOUT`` — + would mute the agent entirely against an SDK that emits a single + undifferentiated stream, turning a cosmetic mismatch into a total + transport failure. + """ + channel = getattr(frame, "channel", None) + name = str(getattr(channel, "name", None) or channel or "").upper() + return not any(side in name for side in _SIDE_CHANNELS) + + def _log_non_stdout(self, frame: Any) -> None: + payload = frame.payload + if isinstance(payload, bytes): + payload = payload.decode(errors="replace") + if payload and payload.strip(): + logger.debug( + "AgentCore shell %s: %s", + getattr(frame.channel, "name", frame.channel), + payload.strip()[:500], + ) async def _write_env_file(self, env: dict[str, str]) -> str: """Materialize *env* as a mode-0600 file inside the session container. @@ -223,19 +269,25 @@ def _closed(msg: str, diagnosis: str) -> TransportClosedError: raise _closed("AgentCore shell closed", "pty_error") timeout = _readline_timeout_sec() try: - return await asyncio.wait_for(self._line_buffer.get(), timeout=timeout) + line = await asyncio.wait_for(self._line_buffer.get(), timeout=timeout) except TimeoutError as e: - # A dead reader task means the WebSocket dropped; report that rather - # than the timeout, which would misattribute an infra failure to a - # silent agent. + raise _closed( + f"AgentCore shell readline timeout ({timeout:g}s)", "pty_error" + ) from e + if line is _READER_ENDED: + # The reader finished — either an error or a clean EOF. Both mean + # the transport is gone, and both must surface now rather than at + # the read timeout. if self._failure is not None: raise _closed( f"AgentCore shell transport failed: {self._failure}", "remote_session_killed", ) from self._failure raise _closed( - f"AgentCore shell readline timeout ({timeout:g}s)", "pty_error" - ) from e + "AgentCore shell closed by the remote session", + "remote_session_killed", + ) + return line async def writeline(self, data: str) -> None: if not self._shell or self._closed: @@ -257,4 +309,10 @@ async def close(self) -> None: @property def is_running(self) -> bool: - return self._shell is not None and not self._closed + """Liveness includes the reader: a dead reader is a dead transport.""" + return ( + self._shell is not None + and not self._closed + and not self._reader_done + and self._failure is None + ) diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py index 2902ac1da..3d772c874 100644 --- a/tests/test_agentcore_builder.py +++ b/tests/test_agentcore_builder.py @@ -127,17 +127,29 @@ def test_symlinks_are_not_packaged(self, tmp_path): def test_buildspec_enforces_the_image_size_cap_remotely(self): """The 2 GB cap must be caught on the worker, before the push.""" - from benchflow.sandbox import agentcore_provisioning as provisioning - commands = " ".join(builders._BUILDSPEC["phases"]["build"]["commands"]) - assert str(provisioning.MAX_IMAGE_MB) in commands assert "BENCHFLOW_IMAGE_TOO_LARGE" in commands # The push must come after the gate, or an oversized image still lands. assert commands.index("BENCHFLOW_IMAGE_TOO_LARGE") < commands.index( "docker push" ) + def test_buildspec_gate_compares_bytes_not_floored_megabytes(self): + """Cap + 1 byte floors to the cap and would slip a megabyte compare. + + The remote gate must reject exactly what image_size_error() rejects. + """ + from benchflow.sandbox import agentcore_provisioning as provisioning + + cap_bytes = provisioning.MAX_IMAGE_MB * 1024 * 1024 + commands = " ".join(builders._BUILDSPEC["phases"]["build"]["commands"]) + + assert f'"$SIZE" -gt {cap_bytes}' in commands + # Local and remote gates must agree on the boundary. + assert provisioning.image_size_error(cap_bytes, "img") is None + assert provisioning.image_size_error(cap_bytes + 1, "img") is not None + def test_buildspec_builds_arm64(self): """AgentCore runs arm64 only; an x86 image would fail opaquely.""" commands = " ".join(builders._BUILDSPEC["phases"]["build"]["commands"]) diff --git a/tests/test_agentcore_parallel.py b/tests/test_agentcore_parallel.py index 1baf98975..9488ff2b1 100644 --- a/tests/test_agentcore_parallel.py +++ b/tests/test_agentcore_parallel.py @@ -495,6 +495,7 @@ async def test_a_runtime_bound_to_another_image_is_updated( "agentRuntimeArtifact": { "containerConfiguration": {"containerUri": "new.example/repo@sha256:b"} }, + "lifecycleConfiguration": env._lifecycle_configuration(), } with ( @@ -526,14 +527,92 @@ def test_adoption_fails_closed_when_the_image_still_mismatches(self): control.get_agent_runtime.return_value = { "agentRuntimeArtifact": { "containerConfiguration": {"containerUri": "old.example/repo@sha256:a"} - } + }, + "lifecycleConfiguration": { + "idleRuntimeSessionTimeout": 900, + "maxLifetime": 28800, + }, } with pytest.raises(SandboxStartupError, match="wrong image"): - AgentCoreSandbox._verify_runtime_image( - control, "rt-1", "new.example/repo@sha256:b" + AgentCoreSandbox._verify_adopted_runtime( + control, + "rt-1", + "new.example/repo@sha256:b", + {"idleRuntimeSessionTimeout": 900, "maxLifetime": 28800}, + ) + + def test_adoption_fails_closed_when_lifecycle_drifted(self): + """An adopted runtime on service defaults reclaims sessions early. + + Live-observed: an update that omits lifecycleConfiguration silently + resets a configured 600/7200 window to the 900/28800 defaults. + """ + from benchflow.sandbox.agentcore import AgentCoreSandbox + from benchflow.sandbox.protocol import SandboxStartupError + + control = MagicMock() + control.get_agent_runtime.return_value = { + "agentRuntimeArtifact": { + "containerConfiguration": {"containerUri": "repo@sha256:a"} + }, + "lifecycleConfiguration": { + "idleRuntimeSessionTimeout": 900, + "maxLifetime": 28800, + }, + } + + with pytest.raises(SandboxStartupError, match="lifecycle"): + AgentCoreSandbox._verify_adopted_runtime( + control, + "rt-1", + "repo@sha256:a", + {"idleRuntimeSessionTimeout": 600, "maxLifetime": 7200}, ) + @pytest.mark.asyncio + async def test_update_preserves_the_configured_lifecycle( + self, tmp_path, monkeypatch + ): + """The rebind must carry lifecycle, or AWS resets it to defaults.""" + from botocore.exceptions import ClientError + + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + monkeypatch.setenv("BENCHFLOW_AGENTCORE_IDLE_TIMEOUT_SEC", "600") + monkeypatch.setenv("BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC", "7200") + env = _sandbox(_make_task(tmp_path)) + control = MagicMock() + control.create_agent_runtime.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "exists"}}, + "CreateAgentRuntime", + ) + control.get_agent_runtime.return_value = { + "status": "READY", + "agentRuntimeArtifact": { + "containerConfiguration": {"containerUri": "repo@sha256:new"} + }, + "lifecycleConfiguration": { + "idleRuntimeSessionTimeout": 600, + "maxLifetime": 7200, + }, + } + + with ( + patch.object(env, "_client", return_value=control), + patch.object( + provisioning, + "find_runtime_by_name", + return_value=("arn:rt", "rt-1", "repo@sha256:old"), + ), + ): + await env._create_or_adopt_runtime("bf_x", "repo@sha256:new") + + sent = control.update_agent_runtime.call_args.kwargs + assert sent["lifecycleConfiguration"] == { + "idleRuntimeSessionTimeout": 600, + "maxLifetime": 7200, + } + class TestComposeRejection: def test_compose_task_is_refused_at_construction(self, tmp_path): @@ -572,3 +651,112 @@ def test_docker_still_accepts_compose_tasks(self, tmp_path): ) assert not any("compose" in issue.reason for issue in issues) + + +class TestLeaseProtectsActiveRuntimes: + """Runtime age is not a session-activity signal. + + Session traffic does not move a runtime's ``lastUpdatedAt``, and there is + no API that enumerates a runtime's active sessions (``ListSessions`` is + Memory-scoped). An old runtime serving a matrix right now is therefore + indistinguishable from an idle one by age alone — which is how cleanup + selected a runtime whose session was mid-command. The lease is the explicit + contract that closes that gap. + """ + + def _control(self, tags): + control = MagicMock() + paginator = MagicMock() + paginator.paginate.return_value = [ + { + "agentRuntimes": [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + # Deliberately ancient: age alone would select it. + "lastUpdatedAt": datetime.now(UTC) - timedelta(days=30), + } + ] + } + ] + control.get_paginator.return_value = paginator + control.list_tags_for_resource.return_value = {"tags": tags} + return control + + def _managed(self, **extra): + return {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, **extra} + + def test_an_old_but_leased_runtime_is_not_deleted(self): + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + future = (datetime.now(UTC) + timedelta(hours=4)).isoformat() + control = self._control(self._managed(**{provisioning.LEASE_TAG: future})) + + report = reap_stale_runtimes(control, max_age_minutes=0) + + assert report.deleted == [] + assert report.skipped_active == 1 + control.delete_agent_runtime.assert_not_called() + + def test_an_expired_lease_allows_reaping(self): + """Positive control: the lease must not block cleanup forever.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + past = (datetime.now(UTC) - timedelta(hours=4)).isoformat() + control = self._control(self._managed(**{provisioning.LEASE_TAG: past})) + + report = reap_stale_runtimes(control, max_age_minutes=0) + + assert report.deleted == ["mine"] + + def test_an_unparseable_lease_is_treated_as_active(self): + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control(self._managed(**{provisioning.LEASE_TAG: "soon"})) + + report = reap_stale_runtimes(control, max_age_minutes=0) + + assert report.deleted == [] + assert report.skipped_active == 1 + + def test_unreadable_tags_are_never_deleted(self): + """No tags means no proof of anything, including that it is ours.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control({}) + control.list_tags_for_resource.side_effect = RuntimeError("denied") + + report = reap_stale_runtimes(control, max_age_minutes=0) + + assert report.deleted == [] + assert report.skipped_unmanaged == 1 + + def test_negative_age_is_rejected(self): + """A sign mistake must not reach delete_agent_runtime.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control(self._managed()) + + with pytest.raises(ValueError, match="must be >= 0"): + reap_stale_runtimes(control, max_age_minutes=-1) + + control.delete_agent_runtime.assert_not_called() + + @pytest.mark.asyncio + async def test_provisioning_writes_a_lease(self, tmp_path, monkeypatch): + """Without this the reaper has no activity signal at all.""" + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + env = _sandbox(_make_task(tmp_path)) + control = MagicMock() + control.create_agent_runtime.return_value = { + "agentRuntimeId": "rt-1", + "agentRuntimeArn": "arn:rt", + } + control.get_agent_runtime.return_value = {"status": "READY"} + + with patch.object(env, "_client", return_value=control): + await env._create_or_adopt_runtime("bf_x", "repo@sha256:a") + + tags = control.tag_resource.call_args.kwargs["tags"] + assert provisioning.LEASE_TAG in tags + assert datetime.fromisoformat(tags[provisioning.LEASE_TAG]) > datetime.now(UTC) diff --git a/tests/test_agentcore_transport.py b/tests/test_agentcore_transport.py new file mode 100644 index 000000000..c9840cf87 --- /dev/null +++ b/tests/test_agentcore_transport.py @@ -0,0 +1,143 @@ +"""AgentCore shell transport: liveness and channel separation. + +The transport carries ACP JSON-RPC. Two failure modes matter more than +throughput: a dead reader that nobody notices (the run hangs until a 15-minute +read timeout instead of reporting an infrastructure disconnect), and side +channels leaking into the protocol stream (a diagnostic line that the JSON-RPC +parser sees as traffic). +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from benchflow.diagnostics import TransportClosedError +from benchflow.sandbox.process.agentcore import AgentCoreProcess + + +def _frame(payload: bytes, channel: str = "STDOUT"): + return SimpleNamespace( + payload=payload, channel=SimpleNamespace(name=channel), raw_channel_byte=1 + ) + + +class _FakeShell: + """Async-iterable stand-in for the SDK's ShellSession.""" + + def __init__(self, frames, *, error: Exception | None = None): + self._frames = list(frames) + self._error = error + self.sent: list[str] = [] + + def __aiter__(self): + return self + + async def __anext__(self): + if self._frames: + return self._frames.pop(0) + if self._error: + raise self._error + raise StopAsyncIteration + + async def send(self, data): + self.sent.append(data) + + async def close(self): + return None + + +def _process(shell): + proc = AgentCoreProcess(MagicMock(), "arn:rt", "s" * 40, "us-west-2") + proc._shell = shell + return proc + + +class TestReaderLiveness: + @pytest.mark.asyncio + async def test_reader_error_wakes_readline_immediately(self): + """A disconnect must not wait out the 900s read timeout.""" + proc = _process(_FakeShell([], error=ConnectionResetError("socket died"))) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + with pytest.raises(TransportClosedError) as excinfo: + await asyncio.wait_for(proc.readline(), timeout=2) + + assert "transport failed" in str(excinfo.value) + assert excinfo.value.diagnostic.transport_diagnosis == "remote_session_killed" + + @pytest.mark.asyncio + async def test_clean_eof_also_wakes_readline(self): + """A clean iterator end records no failure but is still fatal.""" + proc = _process(_FakeShell([])) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + with pytest.raises(TransportClosedError) as excinfo: + await asyncio.wait_for(proc.readline(), timeout=2) + + assert "closed by the remote session" in str(excinfo.value) + + @pytest.mark.asyncio + async def test_buffered_output_is_delivered_before_the_end_sentinel(self): + """Ending the reader must not discard already-framed lines.""" + proc = _process(_FakeShell([_frame(b'{"jsonrpc":"2.0"}\n')])) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + line = await asyncio.wait_for(proc.readline(), timeout=2) + + assert line == b'{"jsonrpc":"2.0"}\n' + + @pytest.mark.asyncio + async def test_is_running_is_false_once_the_reader_ends(self): + """Liveness that ignores the reader reports a dead pipe as healthy.""" + proc = _process(_FakeShell([])) + assert proc.is_running is True + + await proc._drain_frames() + + assert proc.is_running is False + + +class TestChannelSeparation: + @pytest.mark.asyncio + async def test_stderr_never_enters_the_acp_stream(self): + """A diagnostic must not be readable as JSON-RPC traffic.""" + proc = _process( + _FakeShell( + [ + _frame(b"traceback: something failed\n", channel="STDERR"), + _frame(b'{"jsonrpc":"2.0","id":1}\n'), + ] + ) + ) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + line = await asyncio.wait_for(proc.readline(), timeout=2) + + assert line == b'{"jsonrpc":"2.0","id":1}\n' + + @pytest.mark.asyncio + async def test_unknown_channels_are_still_treated_as_stdout(self): + """Older SDKs emit one undifferentiated stream; don't mute the agent.""" + proc = _process(_FakeShell([_frame(b"hello\n", channel="")])) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + assert await asyncio.wait_for(proc.readline(), timeout=2) == b"hello\n" + + +class TestEnvHandling: + @pytest.mark.asyncio + async def test_secrets_are_staged_not_typed_into_the_terminal(self): + """The PTY echoes input, so a typed secret lands in the agent log.""" + proc = _process(_FakeShell([])) + proc._sandbox.write_text_file = AsyncMock(return_value=True) + + path = await proc._write_env_file({"API_KEY": "hunter2"}) + + assert path.startswith("/tmp/") + body = proc._sandbox.write_text_file.call_args.args[1] + assert "hunter2" in body + assert proc._sandbox.write_text_file.call_args.kwargs["mode"] == "600" From 2b69b4b622f5055510808a6711d94fea5cc59b58 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 10:47:18 -0700 Subject: [PATCH 06/13] fix(agentcore): close round-3 review findings; guard deprecated cleanup alias MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Seven P1 and two P2 findings, each re-verified live. The regression first: `bench environment cleanup --max-age -1` stayed destructive because the previous round guarded only `bench sandbox cleanup`. Both now carry `min=0`, and the guard also moved into `reap_stale_sandboxes` itself — every age knob rejects negatives before the cutoff is computed, so a sign mistake cannot mark STARTED sandboxes for deletion through any caller. Lease integrity, in three parts. A failed tag write was swallowed, leaving an unleased runtime while sessions ran against it; it now aborts the launch rather than starting work another process may reap. A managed runtime with no lease is treated as active rather than idle — every runtime BenchFlow provisions is leased before any session starts, so an unleased one is unexplained. And because provisioning is memoized, only the first rollout of an image reached the creation path, so a long or staggered matrix inherited a lease that kept aging; every rollout now refreshes it, throttled to a quarter of the window so this costs a few control-plane calls per run rather than one per rollout. `.dockerignore` character classes (`secret[0-9].pem`) were escaped as literals and leaked into the S3 upload. The translator now honors `[...]` including `[!...]`/`[^...]` negation and ranges. Verified live that the class excludes its matches while a non-matching sibling still ships, so the fix is precise rather than over-broad. Image identity was under-specified twice: only the executable bit was hashed, so 0600 and 0644 collided, and field framing let a path imitate another field. Full permission bits are hashed and every field is length-prefixed. Startup could still hang for the read timeout: the marker wait and the pre-agent drain could both consume the reader's end sentinel. The wait now treats it as a terminal transport failure and the drain re-posts it. Adoption accepted a runtime whose role, network mode, or protocol differed as long as image and lifecycle matched — a different sandbox contract entirely (wrong permissions, wrong egress, or a shell that never answers). All five are verified now. P2: local size inspection failed open and pushed an unmeasurable image, which would resurface as an opaque runtime error; it is now fatal. And a typed but unrecognized shell channel could enter the ACP stream — untyped frames are still treated as agent output so an SDK mismatch cannot mute the agent, but a typed channel must be STDOUT. --- src/benchflow/cli/environment.py | 7 +- src/benchflow/sandbox/agentcore.py | 74 +++++- src/benchflow/sandbox/agentcore_builder.py | 16 +- .../sandbox/agentcore_provisioning.py | 102 ++++++-- src/benchflow/sandbox/agentcore_reaper.py | 9 +- src/benchflow/sandbox/daytona_reaper.py | 14 ++ src/benchflow/sandbox/process/agentcore.py | 48 +++- tests/test_agentcore_builder.py | 73 ++++++ tests/test_agentcore_parallel.py | 220 +++++++++++++++++- tests/test_agentcore_transport.py | 50 +++- 10 files changed, 561 insertions(+), 52 deletions(-) diff --git a/src/benchflow/cli/environment.py b/src/benchflow/cli/environment.py index cbd5ad1ab..6e8062a56 100644 --- a/src/benchflow/cli/environment.py +++ b/src/benchflow/cli/environment.py @@ -145,7 +145,12 @@ def environment_cleanup( bool, typer.Option("--dry-run", help="List sandboxes without deleting") ] = False, max_age_minutes: Annotated[ - int, typer.Option("--max-age", help="Delete sandboxes older than N minutes") + int, + typer.Option( + "--max-age", + min=0, + help="Delete sandboxes older than N minutes", + ), ] = 1440, ) -> None: """Deprecated; use `bench sandbox cleanup`.""" diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py index 2507d52db..e7bbbe54c 100644 --- a/src/benchflow/sandbox/agentcore.py +++ b/src/benchflow/sandbox/agentcore.py @@ -50,6 +50,7 @@ import shlex import tarfile import threading +import time import uuid from datetime import UTC, datetime, timedelta from pathlib import Path, PurePosixPath @@ -83,6 +84,8 @@ # POST /invocations), so register the runtime as HTTP rather than relying on # whatever the service currently defaults to. _PROTOCOL_CONFIGURATION = {"serverProtocol": "HTTP"} +#: AgentCore offers PUBLIC or VPC only; see the registry's enforces_no_network. +_NETWORK_CONFIGURATION = {"networkMode": "PUBLIC"} # Environment overrides. _ENV_REGION = "BENCHFLOW_AGENTCORE_REGION" @@ -272,6 +275,10 @@ async def start(self, force_build: bool) -> None: self.runtime_arn = await self._ensure_runtime(image_uri) # AgentCore requires a session id of at least 33 characters. self.runtime_session_id = f"{uuid.uuid4()}-{uuid.uuid4().hex[:8]}" + # Provisioning is memoized, so only the first rollout of an image + # runs the creation path. Every rollout refreshes the lease (rate + # limited) or a long matrix outlives the lease it inherited. + await self._renew_lease() await self._warm_session() self.logger.info( "AgentCore session ready (runtime=%s, session=%s)", @@ -462,7 +469,7 @@ async def _create_or_adopt_runtime( "containerConfiguration": {"containerUri": image_uri} }, roleArn=self._require_role_arn(), - networkConfiguration={"networkMode": "PUBLIC"}, + networkConfiguration=_NETWORK_CONFIGURATION, protocolConfiguration=_PROTOCOL_CONFIGURATION, lifecycleConfiguration=self._lifecycle_configuration(), description=f"BenchFlow sandbox for {self.environment_name}", @@ -501,7 +508,7 @@ async def _create_or_adopt_runtime( "containerConfiguration": {"containerUri": image_uri} }, roleArn=self._require_role_arn(), - networkConfiguration={"networkMode": "PUBLIC"}, + networkConfiguration=_NETWORK_CONFIGURATION, protocolConfiguration=_PROTOCOL_CONFIGURATION, # Omitting this resets idle/max-lifetime to the service # defaults, silently discarding the caller's configured @@ -515,6 +522,9 @@ async def _create_or_adopt_runtime( runtime_id, image_uri, self._lifecycle_configuration(), + self._require_role_arn(), + _NETWORK_CONFIGURATION, + _PROTOCOL_CONFIGURATION, ) self.logger.info("Adopted existing AgentCore runtime %s", arn) return arn, runtime_id @@ -529,6 +539,24 @@ async def _create_or_adopt_runtime( ) return created["agentRuntimeArn"], runtime_id + async def _renew_lease(self) -> None: + """Refresh this runtime's lease if it is due, before the session runs.""" + if not self.runtime_arn: + return + lifecycle = self._lifecycle_configuration() + window = float( + max(lifecycle["maxLifetime"], lifecycle["idleRuntimeSessionTimeout"]) + ) + if not provisioning.lease_needs_renewal( + self.runtime_arn, window, time.monotonic() + ): + return + await asyncio.to_thread( + self._write_lease, + self._client("bedrock-agentcore-control"), + self.runtime_arn, + ) + def _write_lease(self, control: Any, runtime_arn: str) -> None: """Mark the runtime as possibly-in-use until the session window closes. @@ -552,9 +580,15 @@ def _write_lease(self, control: Any, runtime_arn: str) -> None: tags={provisioning.LEASE_TAG: until.isoformat()}, ) except Exception as exc: - # A runtime without a readable lease is treated as in-use by the - # reaper, so a failure here is conservative, not dangerous. - self.logger.warning("Could not write AgentCore lease: %s", exc) + # Swallowing this would leave the runtime unleased while sessions + # run against it, which is precisely the state cleanup is allowed + # to delete. Fail the launch instead of starting work that another + # process may reap out from under us. + raise SandboxStartupError( + f"Could not write the AgentCore lease for {runtime_arn}: {exc}. " + "Refusing to start a session that cleanup could delete.", + sandbox_id=runtime_arn, + ) from exc @staticmethod def _verify_adopted_runtime( @@ -562,6 +596,9 @@ def _verify_adopted_runtime( runtime_id: str, image_uri: str, lifecycle: dict[str, int], + role_arn: str, + network: dict[str, Any], + protocol: dict[str, Any], ) -> None: """Fail closed unless the adopted runtime matches what this run needs. @@ -581,16 +618,35 @@ def _verify_adopted_runtime( sandbox_id=runtime_id, ) actual = detail.get("lifecycleConfiguration") or {} - drifted = { + drifted: dict[str, tuple[Any, Any]] = { key: (actual.get(key), value) for key, value in lifecycle.items() if actual.get(key) != value } + # Everything the create request pins must also hold on adoption. A + # runtime with the right image but the wrong execution role, network + # mode, or protocol is a different sandbox contract than this run asked + # for — wrong permissions, wrong egress, or a shell that never answers. + for label, expected, found in ( + ("roleArn", role_arn, detail.get("roleArn")), + ( + "networkConfiguration", + dict(network), + dict(detail.get("networkConfiguration") or {}), + ), + ( + "protocolConfiguration", + dict(protocol), + dict(detail.get("protocolConfiguration") or {}), + ), + ): + if found != expected: + drifted[label] = (found, expected) if drifted: raise SandboxStartupError( - f"AgentCore runtime {runtime_id} lifecycle is {actual}, not the " - f"configured {lifecycle} (drift: {drifted}). Sessions could be " - "reclaimed earlier than this run expects.", + f"AgentCore runtime {runtime_id} does not match this run's " + f"configuration (found, expected): {drifted}. Refusing to adopt " + "a runtime with a different contract.", sandbox_id=runtime_id, ) diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py index f51f1f0ba..35cc06856 100644 --- a/src/benchflow/sandbox/agentcore_builder.py +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -164,14 +164,22 @@ async def _reject_oversized(self, image_uri: str) -> None: inspect = await asyncio.to_thread( _run, "docker", "image", "inspect", "-f", "{{.Size}}", image_uri ) + from benchflow.sandbox.protocol import SandboxStartupError + raw = (inspect.stdout or "").strip() if inspect.returncode != 0 or not raw.isdigit(): - logger.debug("Could not measure image size for %s", image_uri) - return + # Failing open here pushes an image whose size is unknown; if it is + # over the hard 2 GB cap the failure resurfaces later as an opaque + # runtime error that reads as a task failure. + raise SandboxStartupError( + f"Could not measure the size of {image_uri} " + f"(docker image inspect exited {inspect.returncode}). Refusing " + "to push an image that cannot be checked against AgentCore's " + f"{provisioning.MAX_IMAGE_MB} MB limit.", + sandbox_id=image_uri, + ) message = provisioning.image_size_error(int(raw), image_uri) if message: - from benchflow.sandbox.protocol import SandboxStartupError - raise SandboxStartupError(message, sandbox_id=image_uri) def _login(self, registry: str) -> None: diff --git a/src/benchflow/sandbox/agentcore_provisioning.py b/src/benchflow/sandbox/agentcore_provisioning.py index 5420fc45e..9c14e7b49 100644 --- a/src/benchflow/sandbox/agentcore_provisioning.py +++ b/src/benchflow/sandbox/agentcore_provisioning.py @@ -57,6 +57,8 @@ _GLOBAL_LOCK = asyncio.Lock() _KEY_LOCKS: dict[str, asyncio.Lock] = {} _RESULTS: dict[str, Any] = {} +#: runtime ARN -> monotonic time of the last lease refresh by this process. +_LEASE_RENEWED: dict[str, float] = {} async def once[T](key: str, factory: Callable[[], Awaitable[T]]) -> T: @@ -85,17 +87,40 @@ def reset_cache() -> None: """Drop memoized provisioning state (tests only).""" _RESULTS.clear() _KEY_LOCKS.clear() + _LEASE_RENEWED.clear() + + +def lease_needs_renewal(runtime_arn: str, window_seconds: float, now: float) -> bool: + """Whether this process should refresh *runtime_arn*'s lease. + + Provisioning is memoized, so only the first rollout of an image reaches the + creation path — every later rollout would inherit a lease that keeps aging + while it runs. A long or staggered matrix therefore ends up with active + sessions on an expired lease, which is the deletion hazard the lease exists + to prevent. + + Renewal is throttled to a quarter of the lease window so the refresh costs + a handful of control-plane calls per runtime per run rather than one per + rollout (``TagResource`` shares the tight control-plane budget). + """ + interval = max(window_seconds / 4, 60.0) + last = _LEASE_RENEWED.get(runtime_arn) + if last is not None and now - last < interval: + return False + _LEASE_RENEWED[runtime_arn] = now + return True def _translate_ignore_pattern(pattern: str) -> re.Pattern[str]: """Compile one ``.dockerignore`` pattern to a regex over relative paths. - Follows Docker's matcher rather than approximating it with ``fnmatch``: - leading and trailing separators are stripped (so ``/secret.env`` is the - root file, not an absolute path that never matches), ``*`` and ``?`` stop - at a separator, and ``**`` spans zero or more path segments. Approximating - these let ``/secret.env`` and ``**/*.pem`` through, which uploaded root - secrets to S3. + Follows Docker's matcher rather than approximating it: leading and trailing + separators are stripped (so ``/secret.env`` is the root file, not an + absolute path that never matches), ``*`` and ``?`` stop at a separator, + ``**`` spans zero or more path segments, and ``[...]`` character classes + are honored. Approximating any of these leaks files Docker excludes into + the CodeBuild upload, which is a credential-exposure bug because that + archive goes to S3. """ cleaned = pattern.strip().strip("/") segments = cleaned.split("/") @@ -105,18 +130,44 @@ def _translate_ignore_pattern(pattern: str) -> re.Pattern[str]: if segment == "**": regex += "(?:.*)?" if last else "(?:[^/]+/)*" continue - for char in segment: - if char == "*": - regex += "[^/]*" - elif char == "?": - regex += "[^/]" - else: - regex += re.escape(char) + regex += _translate_segment(segment) if not last: regex += "/" return re.compile(f"^{regex}$") +def _translate_segment(segment: str) -> str: + """Translate one path segment's wildcards, including character classes.""" + out = "" + index = 0 + while index < len(segment): + char = segment[index] + if char == "*": + out += "[^/]*" + elif char == "?": + out += "[^/]" + elif char == "[": + close = segment.find("]", index + 2) + if close == -1: + # An unterminated class is a literal bracket, as in shell + # globbing — not a syntax error. + out += re.escape(char) + else: + body = segment[index + 1 : close] + negate = body[:1] in ("!", "^") + if negate: + body = body[1:] + # Keep ranges intact; escape only what would change meaning. + body = body.replace("\\", "\\\\").replace("]", "\\]") + out += f"[{'^' if negate else ''}{body}]" + index = close + 1 + continue + else: + out += re.escape(char) + index += 1 + return out + + def _dockerignore_matcher(context_dir: Path) -> Callable[[str], bool]: """Compile ``.dockerignore`` into a predicate over context-relative paths. @@ -188,14 +239,25 @@ def build_context_digest( container even though every byte of content is unchanged. """ digest = hashlib.sha256() - digest.update(dockerfile_text.encode()) - digest.update(b"\0shim\0") - digest.update(shim_text.encode()) + + def field(label: bytes, payload: bytes) -> None: + # Length-prefixed fields: without this a file literally named "shim", + # or a path containing the separator byte, could be framed to produce + # the same digest as a different context. + digest.update(label) + digest.update(str(len(payload)).encode()) + digest.update(b":") + digest.update(payload) + + field(b"dockerfile", dockerfile_text.encode()) + field(b"shim", shim_text.encode()) for path, relative in iter_context_files(context_dir): - digest.update(relative.encode()) - digest.update(b"\0") - digest.update(b"x" if path.stat().st_mode & 0o111 else b"-") - digest.update(hashlib.sha256(path.read_bytes()).digest()) + field(b"path", relative.encode()) + # Full permission bits, not just the executable flag: 0600 and 0644 + # produce different containers, and a secret dropped to 0644 during an + # upgrade must not reuse the old image. + field(b"mode", format(path.stat().st_mode & 0o7777, "04o").encode()) + field(b"blob", hashlib.sha256(path.read_bytes()).digest()) return digest.hexdigest() diff --git a/src/benchflow/sandbox/agentcore_reaper.py b/src/benchflow/sandbox/agentcore_reaper.py index f2ec670af..7657b13c2 100644 --- a/src/benchflow/sandbox/agentcore_reaper.py +++ b/src/benchflow/sandbox/agentcore_reaper.py @@ -91,7 +91,14 @@ def _lease_is_active(tags: Mapping[str, str], now: datetime) -> bool: """ raw = tags.get(LEASE_TAG) if not raw: - return False + # Every runtime BenchFlow provisions is leased before any session runs, + # and that write is fatal if it fails. A managed runtime with no lease + # is therefore unexplained, not idle — treat it as in use. + logger.warning( + "BenchFlow-managed AgentCore runtime has no lease tag; " + "treating it as active rather than deleting it" + ) + return True try: until = datetime.fromisoformat(raw) except ValueError: diff --git a/src/benchflow/sandbox/daytona_reaper.py b/src/benchflow/sandbox/daytona_reaper.py index 2fe90ee75..6354ae033 100644 --- a/src/benchflow/sandbox/daytona_reaper.py +++ b/src/benchflow/sandbox/daytona_reaper.py @@ -172,6 +172,20 @@ def reap_stale_sandboxes( ``{"found", "deleted", "skipped", "failed"}`` (``found`` counts every sandbox listed; foreign ones fall into ``skipped``). """ + for name, value in ( + ("max_age_minutes", max_age_minutes), + ("failed_max_age_minutes", failed_max_age_minutes), + ("min_idle_minutes", min_idle_minutes), + ): + if value < 0: + # A negative age moves the cutoff into the future, so every + # sandbox — including one that is STARTED and serving a run — + # looks past its TTL and is selected for deletion. + raise ValueError( + f"{name} must be >= 0, got {value}. A negative age would " + "select every sandbox, including running ones." + ) + from datetime import UTC, datetime if client is None: diff --git a/src/benchflow/sandbox/process/agentcore.py b/src/benchflow/sandbox/process/agentcore.py index deee6a00e..16e6b6f0f 100644 --- a/src/benchflow/sandbox/process/agentcore.py +++ b/src/benchflow/sandbox/process/agentcore.py @@ -133,15 +133,22 @@ async def _drain_frames(self) -> None: def _is_stdout(frame: Any) -> bool: """Whether *frame* carries agent stdout rather than a side channel. - Side channels are named explicitly and everything else is treated as - agent output. The inverse — allowing only a recognized ``STDOUT`` — - would mute the agent entirely against an SDK that emits a single - undifferentiated stream, turning a cosmetic mismatch into a total - transport failure. + A frame that carries **no** channel information is treated as agent + output, so an SDK emitting one undifferentiated stream is not muted. + A frame that *is* typed must be ``STDOUT`` — an unrecognized typed + channel is a side channel this code has not learned about yet, and + letting it through would put non-protocol bytes into JSON-RPC input. """ channel = getattr(frame, "channel", None) - name = str(getattr(channel, "name", None) or channel or "").upper() - return not any(side in name for side in _SIDE_CHANNELS) + if channel is None: + return True + name = getattr(channel, "name", None) + if name is None: + name = str(channel) + name = str(name).strip().upper() + if not name: + return True + return "STDOUT" in name def _log_non_stdout(self, frame: Any) -> None: payload = frame.payload @@ -242,14 +249,39 @@ async def _await_marker(self, marker: str) -> None: transport_diagnosis="pty_startup_timeout", ), ) from e + if line is _READER_ENDED: + # The shell died during startup. Without this the sentinel is + # consumed here (or cleared by the drain below) and the next + # readline waits out the full read timeout instead. + msg = ( + "AgentCore shell closed before the start marker " + f"(session={self._session_id}): {self._failure or 'EOF'}" + ) + raise TransportClosedError( + msg, + TransportClosedDiagnostic( + raw_message=msg[:500], + transport_diagnosis="remote_session_killed", + ), + ) from self._failure if marker in line.decode(errors="replace"): return def _clear_buffered_output(self) -> None: + """Drop pre-agent terminal noise, but never the end sentinel. + + Discarding it would strand a later readline for the whole read timeout + even though the transport is already known to be gone. + """ self._partial = b"" + saw_end = False while not self._line_buffer.empty(): with contextlib.suppress(asyncio.QueueEmpty): - self._line_buffer.get_nowait() + if self._line_buffer.get_nowait() is _READER_ENDED: + saw_end = True + if saw_end or self._reader_done: + with contextlib.suppress(Exception): + self._line_buffer.put_nowait(_READER_ENDED) async def readline(self) -> bytes: from benchflow.diagnostics import ( diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py index 3d772c874..ce81ed19a 100644 --- a/tests/test_agentcore_builder.py +++ b/tests/test_agentcore_builder.py @@ -248,3 +248,76 @@ def test_negation_re_includes_a_file(self, tmp_path): assert "a.log" not in relatives assert "keep.log" in relatives + + def test_character_classes_are_honoured(self, tmp_path): + """`secret[0-9].pem` is a valid Docker pattern; missing it leaks to S3.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + for name in ("secret1.pem", "secret9.pem", "secretX.pem"): + (tmp_path / name).write_text("blob") + (tmp_path / ".dockerignore").write_text("secret[0-9].pem\n") + + relatives = {rel for _p, rel in provisioning.iter_context_files(tmp_path)} + + assert "secret1.pem" not in relatives + assert "secret9.pem" not in relatives + # Positive control: the class must not over-match. + assert "secretX.pem" in relatives + + def test_permission_bits_change_image_identity(self, tmp_path): + """0600 and 0644 build different containers from identical bytes.""" + import os + + from benchflow.sandbox import agentcore_provisioning as provisioning + + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + script = tmp_path / "entrypoint.sh" + script.write_text("#!/bin/sh\n") + + os.chmod(script, 0o644) + readable = provisioning.build_context_digest(tmp_path, "DF", "SHIM") + os.chmod(script, 0o600) + private = provisioning.build_context_digest(tmp_path, "DF", "SHIM") + os.chmod(script, 0o755) + executable = provisioning.build_context_digest(tmp_path, "DF", "SHIM") + + assert len({readable, private, executable}) == 3 + + def test_shim_bytes_change_image_identity(self, tmp_path): + """The shim is the image entrypoint; an upgrade must not reuse the old.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + + assert provisioning.build_context_digest( + tmp_path, "DF", "shim-v1" + ) != provisioning.build_context_digest(tmp_path, "DF", "shim-v2") + + def test_identity_fields_cannot_be_confused(self, tmp_path): + """Length-prefixed framing: no path can imitate another field.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + (tmp_path / "shim").write_text("x") + + assert provisioning.build_context_digest( + tmp_path, "DF", "abc" + ) != provisioning.build_context_digest(tmp_path, "DFabc", "") + + @pytest.mark.asyncio + async def test_unmeasurable_image_is_not_pushed(self, tmp_path): + """Failing open pushes an image that may exceed the hard 2 GB cap.""" + from benchflow.sandbox.protocol import SandboxStartupError + + builder = builders.LocalDockerBuilder(MagicMock()) + + with ( + patch.object( + builders, + "_run", + return_value=MagicMock(returncode=1, stdout="", stderr="boom"), + ), + pytest.raises(SandboxStartupError, match="Could not measure"), + ): + await builder._reject_oversized("repo:tag") diff --git a/tests/test_agentcore_parallel.py b/tests/test_agentcore_parallel.py index 9488ff2b1..210072bc8 100644 --- a/tests/test_agentcore_parallel.py +++ b/tests/test_agentcore_parallel.py @@ -237,9 +237,15 @@ def test_only_benchflow_managed_runtimes_are_reaped(self): "createdAt": old, }, ] + expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() control = self._control( runtimes, - {"arn:mine": {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}}, + { + "arn:mine": { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expired, + } + }, ) report = reap_stale_runtimes(control, max_age_minutes=60) @@ -258,9 +264,15 @@ def test_recent_runtimes_are_kept(self): "createdAt": datetime.now(UTC) - timedelta(minutes=5), } ] + expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() control = self._control( runtimes, - {"arn:mine": {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}}, + { + "arn:mine": { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expired, + } + }, ) report = reap_stale_runtimes(control, max_age_minutes=1440) @@ -279,9 +291,15 @@ def test_dry_run_reports_without_deleting(self): "createdAt": datetime.now(UTC) - timedelta(days=3), } ] + expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() control = self._control( runtimes, - {"arn:mine": {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}}, + { + "arn:mine": { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expired, + } + }, ) report = reap_stale_runtimes(control, max_age_minutes=60, dry_run=True) @@ -412,8 +430,15 @@ def _control(self, runtimes, tags): } return control - def _managed(self, arn): - return {arn: {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE}} + def _managed(self, arn, *, leased_until=None): + """Managed tags. Default lease is expired, i.e. reapable.""" + expiry = leased_until or (datetime.now(UTC) - timedelta(hours=1)) + return { + arn: { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expiry.isoformat(), + } + } def test_fresh_runtime_in_the_real_list_shape_is_kept(self): """Reproduces the reported P1 with the shape AWS actually returns. @@ -496,6 +521,9 @@ async def test_a_runtime_bound_to_another_image_is_updated( "containerConfiguration": {"containerUri": "new.example/repo@sha256:b"} }, "lifecycleConfiguration": env._lifecycle_configuration(), + "roleArn": "arn:aws:iam::1:role/rt", + "networkConfiguration": {"networkMode": "PUBLIC"}, + "protocolConfiguration": {"serverProtocol": "HTTP"}, } with ( @@ -532,6 +560,9 @@ def test_adoption_fails_closed_when_the_image_still_mismatches(self): "idleRuntimeSessionTimeout": 900, "maxLifetime": 28800, }, + "roleArn": "arn:aws:iam::1:role/rt", + "networkConfiguration": {"networkMode": "PUBLIC"}, + "protocolConfiguration": {"serverProtocol": "HTTP"}, } with pytest.raises(SandboxStartupError, match="wrong image"): @@ -540,6 +571,9 @@ def test_adoption_fails_closed_when_the_image_still_mismatches(self): "rt-1", "new.example/repo@sha256:b", {"idleRuntimeSessionTimeout": 900, "maxLifetime": 28800}, + "arn:aws:iam::1:role/rt", + {"networkMode": "PUBLIC"}, + {"serverProtocol": "HTTP"}, ) def test_adoption_fails_closed_when_lifecycle_drifted(self): @@ -562,12 +596,15 @@ def test_adoption_fails_closed_when_lifecycle_drifted(self): }, } - with pytest.raises(SandboxStartupError, match="lifecycle"): + with pytest.raises(SandboxStartupError, match="does not match"): AgentCoreSandbox._verify_adopted_runtime( control, "rt-1", "repo@sha256:a", {"idleRuntimeSessionTimeout": 600, "maxLifetime": 7200}, + "arn:aws:iam::1:role/rt", + {"networkMode": "PUBLIC"}, + {"serverProtocol": "HTTP"}, ) @pytest.mark.asyncio @@ -595,6 +632,9 @@ async def test_update_preserves_the_configured_lifecycle( "idleRuntimeSessionTimeout": 600, "maxLifetime": 7200, }, + "roleArn": "arn:aws:iam::1:role/rt", + "networkConfiguration": {"networkMode": "PUBLIC"}, + "protocolConfiguration": {"serverProtocol": "HTTP"}, } with ( @@ -760,3 +800,171 @@ async def test_provisioning_writes_a_lease(self, tmp_path, monkeypatch): tags = control.tag_resource.call_args.kwargs["tags"] assert provisioning.LEASE_TAG in tags assert datetime.fromisoformat(tags[provisioning.LEASE_TAG]) > datetime.now(UTC) + + +class TestLeaseIntegrity: + """A lease is only protection if it is written, kept, and refreshed.""" + + @pytest.mark.asyncio + async def test_a_failed_lease_write_aborts_the_launch(self, tmp_path, monkeypatch): + """Swallowing it starts a session on a runtime cleanup may delete.""" + from benchflow.sandbox.protocol import SandboxStartupError + + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + env = _sandbox(_make_task(tmp_path)) + control = MagicMock() + control.create_agent_runtime.return_value = { + "agentRuntimeId": "rt-1", + "agentRuntimeArn": "arn:rt", + } + control.tag_resource.side_effect = RuntimeError("AccessDenied") + + with ( + patch.object(env, "_client", return_value=control), + pytest.raises(SandboxStartupError, match="lease"), + ): + await env._create_or_adopt_runtime("bf_x", "repo@sha256:a") + + def test_a_managed_runtime_without_a_lease_is_not_deleted(self): + """Every provisioned runtime is leased, so an unleased one is unexplained.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = MagicMock() + paginator = MagicMock() + paginator.paginate.return_value = [ + { + "agentRuntimes": [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "lastUpdatedAt": datetime.now(UTC) - timedelta(days=30), + } + ] + } + ] + control.get_paginator.return_value = paginator + control.list_tags_for_resource.return_value = { + "tags": {provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE} + } + + report = reap_stale_runtimes(control, max_age_minutes=0) + + assert report.deleted == [] + assert report.skipped_active == 1 + + def test_renewal_is_due_again_after_the_throttle_window(self): + """Cache-hit rollouts must refresh, or a long matrix outlives its lease.""" + window = 7200.0 + arn = "arn:renew" + + assert provisioning.lease_needs_renewal(arn, window, 0.0) is True + assert provisioning.lease_needs_renewal(arn, window, 10.0) is False + assert provisioning.lease_needs_renewal(arn, window, window) is True + + @pytest.mark.asyncio + async def test_every_rollout_renews_when_due(self, tmp_path, monkeypatch): + """Provisioning is memoized, so renewal cannot live only in creation.""" + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + env = _sandbox(_make_task(tmp_path)) + env.runtime_arn = "arn:rt-renew-probe" + control = MagicMock() + + with patch.object(env, "_client", return_value=control): + await env._renew_lease() + + control.tag_resource.assert_called_once() + assert provisioning.LEASE_TAG in control.tag_resource.call_args.kwargs["tags"] + + +class TestAdoptionContract: + def _detail(self, **overrides): + detail = { + "agentRuntimeArtifact": { + "containerConfiguration": {"containerUri": "repo@sha256:a"} + }, + "lifecycleConfiguration": { + "idleRuntimeSessionTimeout": 600, + "maxLifetime": 7200, + }, + "roleArn": "arn:aws:iam::1:role/rt", + "networkConfiguration": {"networkMode": "PUBLIC"}, + "protocolConfiguration": {"serverProtocol": "HTTP"}, + } + detail.update(overrides) + return detail + + def _verify(self, detail): + from benchflow.sandbox.agentcore import AgentCoreSandbox + + control = MagicMock() + control.get_agent_runtime.return_value = detail + AgentCoreSandbox._verify_adopted_runtime( + control, + "rt-1", + "repo@sha256:a", + {"idleRuntimeSessionTimeout": 600, "maxLifetime": 7200}, + "arn:aws:iam::1:role/rt", + {"networkMode": "PUBLIC"}, + {"serverProtocol": "HTTP"}, + ) + + def test_a_fully_matching_runtime_is_accepted(self): + self._verify(self._detail()) + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("roleArn", "arn:aws:iam::1:role/someone-else"), + ("networkConfiguration", {"networkMode": "VPC"}), + ("protocolConfiguration", {"serverProtocol": "MCP"}), + ], + ) + def test_a_differing_contract_is_refused(self, field, value): + """Right image, wrong contract: wrong permissions, egress, or shell.""" + from benchflow.sandbox.protocol import SandboxStartupError + + with pytest.raises(SandboxStartupError, match="does not match"): + self._verify(self._detail(**{field: value})) + + +class TestDeprecatedCleanupAliasIsSafe: + """`bench environment cleanup` reaches the same destructive code. + + Guarding only the new command left the deprecated alias able to select + STARTED sandboxes for deletion via a negative age. + """ + + def test_negative_age_is_rejected_by_the_daytona_reaper(self): + from benchflow.sandbox.daytona_reaper import reap_stale_sandboxes + + client = MagicMock() + + with pytest.raises(ValueError, match="must be >= 0"): + reap_stale_sandboxes(client, max_age_minutes=-1, dry_run=True) + + client.delete.assert_not_called() + + @pytest.mark.parametrize( + "kwargs", + [ + {"max_age_minutes": -1}, + {"failed_max_age_minutes": -1}, + {"min_idle_minutes": -1}, + ], + ) + def test_every_age_knob_rejects_negatives(self, kwargs): + from benchflow.sandbox.daytona_reaper import reap_stale_sandboxes + + with pytest.raises(ValueError, match="must be >= 0"): + reap_stale_sandboxes(MagicMock(), dry_run=True, **kwargs) + + @pytest.mark.parametrize("command", ["sandbox", "environment"]) + def test_cli_rejects_negative_max_age(self, command): + """Both the current command and the deprecated alias must refuse it.""" + from typer.testing import CliRunner + + from benchflow.cli.main import app + + result = CliRunner().invoke(app, [command, "cleanup", "--max-age", "-1"]) + + assert result.exit_code != 0 diff --git a/tests/test_agentcore_transport.py b/tests/test_agentcore_transport.py index c9840cf87..f510113fb 100644 --- a/tests/test_agentcore_transport.py +++ b/tests/test_agentcore_transport.py @@ -120,13 +120,57 @@ async def test_stderr_never_enters_the_acp_stream(self): assert line == b'{"jsonrpc":"2.0","id":1}\n' @pytest.mark.asyncio - async def test_unknown_channels_are_still_treated_as_stdout(self): - """Older SDKs emit one undifferentiated stream; don't mute the agent.""" - proc = _process(_FakeShell([_frame(b"hello\n", channel="")])) + async def test_frames_without_channel_info_are_treated_as_stdout(self): + """An SDK emitting one undifferentiated stream must not be muted. + + Requiring a recognized STDOUT would turn a naming mismatch into total + transport failure rather than a cosmetic one. + """ + untyped = SimpleNamespace(payload=b"hello\n", channel=None) + proc = _process(_FakeShell([untyped])) proc._reader_task = asyncio.create_task(proc._drain_frames()) assert await asyncio.wait_for(proc.readline(), timeout=2) == b"hello\n" + @pytest.mark.asyncio + async def test_a_typed_unknown_channel_never_reaches_the_acp_stream(self): + """A typed channel this code does not know is still a side channel. + + Admitting it would put non-protocol bytes into JSON-RPC input, which is + how a diagnostic gets parsed as — or impersonates — protocol traffic. + """ + proc = _process( + _FakeShell( + [ + _frame(b"telemetry blob\n", channel="METRICS"), + _frame(b'{"jsonrpc":"2.0","id":7}\n'), + ] + ) + ) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + assert ( + await asyncio.wait_for(proc.readline(), timeout=2) + == b'{"jsonrpc":"2.0","id":7}\n' + ) + + @pytest.mark.asyncio + async def test_shell_death_during_startup_does_not_hang(self): + """Startup drain must not swallow the end sentinel. + + Discarding it strands the next readline for the full 900s timeout even + though the transport is already known to be gone. + """ + proc = _process(_FakeShell([_frame(b"noise\n")])) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + await asyncio.sleep(0) + while not proc._reader_done: + await asyncio.sleep(0.01) + proc._clear_buffered_output() + + with pytest.raises(TransportClosedError): + await asyncio.wait_for(proc.readline(), timeout=2) + class TestEnvHandling: @pytest.mark.asyncio From 4ede269c43b5480657867c03b5db9ded8ab2f4db Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 11:44:02 -0700 Subject: [PATCH 07/13] fix(agentcore): close round-4 findings; guard reserved paths and half-close MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two P1, four P2, one P3 — each mutation-verified rather than merely passing. P1: a failed lease renewal was throttled as if it had succeeded. `lease_needs_renewal()` recorded the timestamp before `TagResource` ran, so a renewal that raised still advanced the window and the immediate retry made zero calls — continuing against a lease that was never extended. The predicate is now pure and `mark_lease_renewed()` is called only after the write lands. P1: reserved generated paths could destroy a task's own files. `materialized()` writes `Dockerfile.benchflow-agentcore` and `.benchflow_agentcore_shim.py` into the caller-owned environment directory and unlinks them afterwards, so a task shipping either name lost it permanently — and because the canonical context walk skips both names, the collision also cache-hit as though the task file did not exist. Both paths are now refused before identity or build, with the originals untouched. P2: typed channel filtering used substring membership, so `NOT_STDOUT` and `STDOUT_METADATA` would have been admitted into the ACP stream. The final enum-qualified component is now compared exactly to `STDOUT`. P2: writes were accepted after reader EOF. `ContainerTransport.send()` does not pre-check liveness, so an ACP request could go to a socket whose read side was already gone and its reply could never arrive — a stall rather than a failure. `writeline()` now gates on the same liveness state as `readline()`. P2: failed-build diagnostics ran on the event loop. `_build_failure()` fetches CloudWatch logs synchronously, so one failed build stalled every concurrent provisioning coroutine; it is dispatched through `asyncio.to_thread` now. P2: a failed delete of an uploaded build context was logged at DEBUG. The archive can carry task source and credentials, and while bucket hardening and the one-day lifecycle bound the exposure, a retained context was invisible for that whole day. It warns with the key and error while preserving the build result. P3: new regression docstrings now name PR #937 per AGENTS.md, and the mutation-blind cases are fixed. The framing test previously passed against the superseded scheme; it now pins an input pair that genuinely collided under it (a shim containing the old `\0shim\0` separator re-splits as a different dockerfile/shim pair) and asserts that collision explicitly before asserting the new framing separates them. Added fail-then-retry, `start()` wiring, and reserved-path byte-preservation regressions. --- src/benchflow/sandbox/agentcore.py | 24 ++- src/benchflow/sandbox/agentcore_builder.py | 14 +- .../sandbox/agentcore_provisioning.py | 16 +- src/benchflow/sandbox/process/agentcore.py | 16 +- tests/test_agentcore_builder.py | 121 +++++++++++++-- tests/test_agentcore_parallel.py | 140 +++++++++++++++++- tests/test_agentcore_transport.py | 72 +++++++++ 7 files changed, 374 insertions(+), 29 deletions(-) diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py index e7bbbe54c..7c7197f8c 100644 --- a/src/benchflow/sandbox/agentcore.py +++ b/src/benchflow/sandbox/agentcore.py @@ -224,6 +224,21 @@ def _validate_definition(self) -> None: f"No Dockerfile found in {self.environment_dir} and no " "docker_image specified in task config." ) + for reserved in ( + provisioning.GENERATED_DOCKERFILE, + provisioning.GENERATED_SHIM, + ): + if (self.environment_dir / reserved).exists(): + # BenchFlow writes these into the caller's own task directory + # and removes them afterwards. A task shipping a file at either + # path would be overwritten and then permanently deleted, and + # the canonical context walk skips both names, so the collision + # would also cache-hit as if the task file did not exist. + raise ValueError( + f"{reserved} already exists in {self.environment_dir}. " + "That path is reserved by the AgentCore backend, which " + "would overwrite and then delete it. Rename the task file." + ) compose = compose_definition_path(self.environment_dir) if compose is not None: # AgentCore builds and runs exactly one container. Accepting a @@ -547,15 +562,18 @@ async def _renew_lease(self) -> None: window = float( max(lifecycle["maxLifetime"], lifecycle["idleRuntimeSessionTimeout"]) ) - if not provisioning.lease_needs_renewal( - self.runtime_arn, window, time.monotonic() - ): + now = time.monotonic() + if not provisioning.lease_needs_renewal(self.runtime_arn, window, now): return await asyncio.to_thread( self._write_lease, self._client("bedrock-agentcore-control"), self.runtime_arn, ) + # Only now — a failed write must leave the throttle untouched so the + # next rollout retries instead of running on a lease that was never + # extended. + provisioning.mark_lease_renewed(self.runtime_arn, now) def _write_lease(self, control: Any, runtime_arn: str) -> None: """Mark the runtime as possibly-in-use until the session window closes. diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py index 35cc06856..727cee292 100644 --- a/src/benchflow/sandbox/agentcore_builder.py +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -281,7 +281,17 @@ async def build_and_push(self, request: BuildRequest) -> None: self._client("s3").delete_object, Bucket=self._bucket, Key=key ) except Exception as exc: - logger.debug("Could not delete build context %s: %s", key, exc) + # The archive can carry task source and credentials. Bucket + # hardening and the one-day lifecycle bound the exposure, so + # this must not replace a successful build result — but at + # DEBUG a retained context is invisible for that whole day. + logger.warning( + "Could not delete uploaded build context s3://%s/%s: %s. " + "It will expire with the bucket lifecycle policy.", + self._bucket, + key, + exc, + ) def _package(self, request: BuildRequest) -> bytes: """Zip the build context with the generated Dockerfile and shim inside. @@ -424,7 +434,7 @@ async def _run_build(self, request: BuildRequest, key: str) -> None: if build["buildStatus"] == "SUCCEEDED": logger.info("Pushed AgentCore image %s (codebuild)", request.image_uri) return - raise self._build_failure(build, request.image_uri) + raise await asyncio.to_thread(self._build_failure, build, request.image_uri) raise TimeoutError( f"CodeBuild {build_id} did not finish within the build timeout" ) diff --git a/src/benchflow/sandbox/agentcore_provisioning.py b/src/benchflow/sandbox/agentcore_provisioning.py index 9c14e7b49..be32c5bd0 100644 --- a/src/benchflow/sandbox/agentcore_provisioning.py +++ b/src/benchflow/sandbox/agentcore_provisioning.py @@ -102,13 +102,21 @@ def lease_needs_renewal(runtime_arn: str, window_seconds: float, now: float) -> Renewal is throttled to a quarter of the lease window so the refresh costs a handful of control-plane calls per runtime per run rather than one per rollout (``TagResource`` shares the tight control-plane budget). + + Pure predicate: it records nothing. The throttle is only advanced by + :func:`mark_lease_renewed` after a write actually lands, so a failed + renewal cannot be remembered as a successful one and suppress the retry + that would have fixed it. """ - interval = max(window_seconds / 4, 60.0) last = _LEASE_RENEWED.get(runtime_arn) - if last is not None and now - last < interval: - return False + if last is None: + return True + return (now - last) >= max(window_seconds / 4, 60.0) + + +def mark_lease_renewed(runtime_arn: str, now: float) -> None: + """Record a lease write that succeeded.""" _LEASE_RENEWED[runtime_arn] = now - return True def _translate_ignore_pattern(pattern: str) -> re.Pattern[str]: diff --git a/src/benchflow/sandbox/process/agentcore.py b/src/benchflow/sandbox/process/agentcore.py index 16e6b6f0f..fcaca911e 100644 --- a/src/benchflow/sandbox/process/agentcore.py +++ b/src/benchflow/sandbox/process/agentcore.py @@ -148,7 +148,11 @@ def _is_stdout(frame: Any) -> bool: name = str(name).strip().upper() if not name: return True - return "STDOUT" in name + # Enum reprs arrive as "SHELLCHANNEL.STDOUT"; compare the final + # component *exactly*. Substring membership admitted NOT_STDOUT and + # STDOUT_METADATA, which contradicts the fail-closed rule even though + # today's SDK happens not to define such names. + return name.rsplit(".", 1)[-1] == "STDOUT" def _log_non_stdout(self, frame: Any) -> None: payload = frame.payload @@ -322,8 +326,14 @@ def _closed(msg: str, diagnosis: str) -> TransportClosedError: return line async def writeline(self, data: str) -> None: - if not self._shell or self._closed: - raise RuntimeError("AgentCore shell not started") + if not self.is_running: + # Half-close guard: ContainerTransport.send() does not pre-check + # liveness, so without this an ACP request is handed to a socket + # whose read side is already gone and the reply can never arrive. + raise RuntimeError( + "AgentCore shell is not running (reader ended); refusing to " + "send to a half-closed transport" + ) await self._shell.send(data + "\n") async def close(self) -> None: diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py index ce81ed19a..328c0628f 100644 --- a/tests/test_agentcore_builder.py +++ b/tests/test_agentcore_builder.py @@ -8,9 +8,11 @@ from __future__ import annotations +import asyncio +import time import zipfile from io import BytesIO -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -250,7 +252,10 @@ def test_negation_re_includes_a_file(self, tmp_path): assert "keep.log" in relatives def test_character_classes_are_honoured(self, tmp_path): - """`secret[0-9].pem` is a valid Docker pattern; missing it leaks to S3.""" + """`secret[0-9].pem` is a valid Docker pattern; missing it leaks to S3. + + Guards PR #937. + """ from benchflow.sandbox import agentcore_provisioning as provisioning (tmp_path / "Dockerfile").write_text("FROM scratch\n") @@ -266,7 +271,10 @@ def test_character_classes_are_honoured(self, tmp_path): assert "secretX.pem" in relatives def test_permission_bits_change_image_identity(self, tmp_path): - """0600 and 0644 build different containers from identical bytes.""" + """0600 and 0644 build different containers from identical bytes. + + Guards PR #937. + """ import os from benchflow.sandbox import agentcore_provisioning as provisioning @@ -285,7 +293,10 @@ def test_permission_bits_change_image_identity(self, tmp_path): assert len({readable, private, executable}) == 3 def test_shim_bytes_change_image_identity(self, tmp_path): - """The shim is the image entrypoint; an upgrade must not reuse the old.""" + """The shim is the image entrypoint; an upgrade must not reuse the old. + + Guards PR #937. + """ from benchflow.sandbox import agentcore_provisioning as provisioning (tmp_path / "Dockerfile").write_text("FROM scratch\n") @@ -295,19 +306,42 @@ def test_shim_bytes_change_image_identity(self, tmp_path): ) != provisioning.build_context_digest(tmp_path, "DF", "shim-v2") def test_identity_fields_cannot_be_confused(self, tmp_path): - """Length-prefixed framing: no path can imitate another field.""" + """Length-prefixed framing kills a real digest collision. + + Guards PR #937. The superseded scheme concatenated the Dockerfile, a + fixed ``\0shim\0`` separator, and the shim with no lengths, so a shim + containing that separator could be re-split as a different + (dockerfile, shim) pair with an identical digest — two different images + sharing one identity, and so one runtime. + + The first assertion pins that the chosen inputs genuinely collided + before; without it this test could pass against the old scheme too. + """ + import hashlib + from benchflow.sandbox import agentcore_provisioning as provisioning - (tmp_path / "Dockerfile").write_text("FROM scratch\n") - (tmp_path / "shim").write_text("x") + def superseded_framing(dockerfile: str, shim: str) -> str: + digest = hashlib.sha256() + digest.update(dockerfile.encode()) + digest.update(b"\0shim\0") + digest.update(shim.encode()) + return digest.hexdigest() + + first = ("A", "B\0shim\0C") + second = ("A\0shim\0B", "C") + assert superseded_framing(*first) == superseded_framing(*second) assert provisioning.build_context_digest( - tmp_path, "DF", "abc" - ) != provisioning.build_context_digest(tmp_path, "DFabc", "") + tmp_path, *first + ) != provisioning.build_context_digest(tmp_path, *second) @pytest.mark.asyncio async def test_unmeasurable_image_is_not_pushed(self, tmp_path): - """Failing open pushes an image that may exceed the hard 2 GB cap.""" + """Failing open pushes an image that may exceed the hard 2 GB cap. + + Guards PR #937. + """ from benchflow.sandbox.protocol import SandboxStartupError builder = builders.LocalDockerBuilder(MagicMock()) @@ -321,3 +355,70 @@ async def test_unmeasurable_image_is_not_pushed(self, tmp_path): pytest.raises(SandboxStartupError, match="Could not measure"), ): await builder._reject_oversized("repo:tag") + + @pytest.mark.asyncio + async def test_failed_context_cleanup_is_warned_not_hidden(self, tmp_path, caplog): + """Guards PR #937: a retained context can hold source and credentials. + + Bucket hardening and the one-day lifecycle bound the exposure, so this + must not fail the build — but at DEBUG the retained object is invisible + for that entire day. + """ + import logging + + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + s3 = MagicMock() + s3.delete_object.side_effect = RuntimeError("AccessDenied") + + with ( + patch.object(builder, "_client", return_value=s3), + patch.object(builder, "_ensure_bucket"), + patch.object(builder, "_ensure_project"), + patch.object(builder, "_run_build", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="benchflow.agentcore-builder"), + ): + await builder.build_and_push(_request(tmp_path)) + + assert any( + "Could not delete uploaded build context" in r.message + for r in caplog.records + ) + + @pytest.mark.asyncio + async def test_failure_diagnostics_do_not_block_the_event_loop(self, tmp_path): + """Guards PR #937: one failed build must not stall other provisioning. + + `_build_failure` fetches CloudWatch logs; running it inline on the loop + froze every concurrent coroutine for the duration of that fetch. + """ + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + codebuild = MagicMock() + codebuild.start_build.return_value = {"build": {"id": "b:1"}} + codebuild.batch_get_builds.return_value = { + "builds": [{"id": "b:1", "buildStatus": "FAILED", "phases": []}] + } + + def _slow_log_tail(_self, _build): + time.sleep(0.25) + return "boom" + + ticks = 0 + + async def _ticker(): + nonlocal ticks + while True: + await asyncio.sleep(0.01) + ticks += 1 + + with ( + patch.object(builder, "_client", return_value=codebuild), + patch.object(builders, "_CODEBUILD_POLL_SEC", 0), + patch.object(builders.CodeBuildBuilder, "_log_tail", _slow_log_tail), + ): + beat = asyncio.create_task(_ticker()) + with pytest.raises(RuntimeError, match="CodeBuild"): + await builder._run_build(_request(tmp_path), "contexts/x.zip") + beat.cancel() + + # A blocking diagnostic starved the ticker to ~1 tick. + assert ticks > 5 diff --git a/tests/test_agentcore_parallel.py b/tests/test_agentcore_parallel.py index 210072bc8..5a1ee2625 100644 --- a/tests/test_agentcore_parallel.py +++ b/tests/test_agentcore_parallel.py @@ -807,7 +807,10 @@ class TestLeaseIntegrity: @pytest.mark.asyncio async def test_a_failed_lease_write_aborts_the_launch(self, tmp_path, monkeypatch): - """Swallowing it starts a session on a runtime cleanup may delete.""" + """Swallowing it starts a session on a runtime cleanup may delete. + + Guards PR #937. + """ from benchflow.sandbox.protocol import SandboxStartupError monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") @@ -826,7 +829,10 @@ async def test_a_failed_lease_write_aborts_the_launch(self, tmp_path, monkeypatc await env._create_or_adopt_runtime("bf_x", "repo@sha256:a") def test_a_managed_runtime_without_a_lease_is_not_deleted(self): - """Every provisioned runtime is leased, so an unleased one is unexplained.""" + """Every provisioned runtime is leased, so an unleased one is unexplained. + + Guards PR #937. + """ from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes control = MagicMock() @@ -853,17 +859,24 @@ def test_a_managed_runtime_without_a_lease_is_not_deleted(self): assert report.skipped_active == 1 def test_renewal_is_due_again_after_the_throttle_window(self): - """Cache-hit rollouts must refresh, or a long matrix outlives its lease.""" + """Cache-hit rollouts must refresh, or a long matrix outlives its lease. + + Guards PR #937. + """ window = 7200.0 arn = "arn:renew" assert provisioning.lease_needs_renewal(arn, window, 0.0) is True + provisioning.mark_lease_renewed(arn, 0.0) assert provisioning.lease_needs_renewal(arn, window, 10.0) is False assert provisioning.lease_needs_renewal(arn, window, window) is True @pytest.mark.asyncio async def test_every_rollout_renews_when_due(self, tmp_path, monkeypatch): - """Provisioning is memoized, so renewal cannot live only in creation.""" + """Provisioning is memoized, so renewal cannot live only in creation. + + Guards PR #937. + """ monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") env = _sandbox(_make_task(tmp_path)) env.runtime_arn = "arn:rt-renew-probe" @@ -920,7 +933,10 @@ def test_a_fully_matching_runtime_is_accepted(self): ], ) def test_a_differing_contract_is_refused(self, field, value): - """Right image, wrong contract: wrong permissions, egress, or shell.""" + """Right image, wrong contract: wrong permissions, egress, or shell. + + Guards PR #937. + """ from benchflow.sandbox.protocol import SandboxStartupError with pytest.raises(SandboxStartupError, match="does not match"): @@ -930,8 +946,8 @@ def test_a_differing_contract_is_refused(self, field, value): class TestDeprecatedCleanupAliasIsSafe: """`bench environment cleanup` reaches the same destructive code. - Guarding only the new command left the deprecated alias able to select - STARTED sandboxes for deletion via a negative age. + Guards PR #937. Guarding only the new command left the deprecated alias + able to select STARTED sandboxes for deletion via a negative age. """ def test_negative_age_is_rejected_by_the_daytona_reaper(self): @@ -968,3 +984,113 @@ def test_cli_rejects_negative_max_age(self, command): result = CliRunner().invoke(app, [command, "cleanup", "--max-age", "-1"]) assert result.exit_code != 0 + + +class TestLeaseRenewalRetry: + """Guards PR #937: a failed renewal must not be throttled as a success.""" + + def _sandbox_with_control(self, tmp_path, control): + env = _sandbox(_make_task(tmp_path)) + env.runtime_arn = "arn:retry-probe" + return env + + @pytest.mark.asyncio + async def test_a_failed_renewal_retries_immediately(self, tmp_path, monkeypatch): + """The throttle recorded the attempt, not the write. + + A first renewal that raised still advanced the window, so the retry + that would have fixed it made zero tag_resource calls and the rollout + continued against the un-extended lease. + """ + from benchflow.sandbox.protocol import SandboxStartupError + + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + env = self._sandbox_with_control(tmp_path, None) + control = MagicMock() + control.tag_resource.side_effect = RuntimeError("AccessDenied") + + with ( + patch.object(env, "_client", return_value=control), + pytest.raises(SandboxStartupError), + ): + await env._renew_lease() + + # AWS recovers; the very next attempt must actually call AWS. + control.tag_resource.side_effect = None + with patch.object(env, "_client", return_value=control): + await env._renew_lease() + + assert control.tag_resource.call_count == 2 + + def test_the_throttle_only_advances_on_success(self): + """Guards PR #937: the predicate must record nothing by itself.""" + arn = "arn:pure-predicate" + window = 7200.0 + + assert provisioning.lease_needs_renewal(arn, window, 0.0) is True + # Still due — checking is not renewing. + assert provisioning.lease_needs_renewal(arn, window, 1.0) is True + + provisioning.mark_lease_renewed(arn, 1.0) + + assert provisioning.lease_needs_renewal(arn, window, 2.0) is False + + @pytest.mark.asyncio + async def test_start_renews_before_warming_the_session(self, tmp_path, monkeypatch): + """Guards PR #937: renewal must be wired into start(), not just callable. + + Provisioning is memoized, so a cache-hit rollout reaches neither the + create nor the adopt path; if start() did not renew, the lease would + only ever be written by the first rollout of an image. + """ + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + env = _sandbox(_make_task(tmp_path)) + order: list[str] = [] + + async def _publish(*, force_build): + return "repo@sha256:a" + + async def _ensure(image_uri): + env.runtime_arn = "arn:started" + return "arn:started" + + async def _renew(): + order.append("renew") + + async def _warm(): + order.append("warm") + + with ( + patch.object(env, "_publish_image", side_effect=_publish), + patch.object(env, "_ensure_runtime", side_effect=_ensure), + patch.object(env, "_renew_lease", side_effect=_renew), + patch.object(env, "_warm_session", side_effect=_warm), + ): + await env.start(force_build=False) + + assert order == ["renew", "warm"] + + +class TestReservedPathsArePreserved: + """Guards PR #937: never clobber a task file at a generated path.""" + + @pytest.mark.parametrize( + "reserved", + ["Dockerfile.benchflow-agentcore", ".benchflow_agentcore_shim.py"], + ) + def test_a_colliding_task_file_is_refused(self, tmp_path, reserved): + """These are written into the caller's tree and then deleted. + + A task shipping either name lost its file permanently, and because the + canonical walk skips both names the collision also cache-hit as though + the task file did not exist. + """ + task_dir = _make_task(tmp_path) + original = "important task content\n" + (task_dir / reserved).write_text(original) + + with pytest.raises(ValueError, match="reserved"): + _sandbox(task_dir) + + # The refusal must happen before anything touches the file. + assert (task_dir / reserved).read_text() == original diff --git a/tests/test_agentcore_transport.py b/tests/test_agentcore_transport.py index f510113fb..d4116591a 100644 --- a/tests/test_agentcore_transport.py +++ b/tests/test_agentcore_transport.py @@ -185,3 +185,75 @@ async def test_secrets_are_staged_not_typed_into_the_terminal(self): body = proc._sandbox.write_text_file.call_args.args[1] assert "hunter2" in body assert proc._sandbox.write_text_file.call_args.kwargs["mode"] == "600" + + +class TestAdversarialChannelNames: + """Guards PR #937: typed channels are matched exactly, not by substring.""" + + @pytest.mark.parametrize("name", ["NOT_STDOUT", "STDOUT_METADATA", "METRICS"]) + def test_colliding_typed_names_are_not_stdout(self, name): + """Substring membership admitted these; only STDOUT may be ACP input.""" + assert AgentCoreProcess._is_stdout(_frame(b"x\n", channel=name)) is False + + @pytest.mark.parametrize("name", ["STDOUT", "ShellChannel.STDOUT", " stdout "]) + def test_genuine_stdout_is_accepted(self, name): + """Enum-qualified and case/space variants must still deliver.""" + assert AgentCoreProcess._is_stdout(_frame(b"x\n", channel=name)) is True + + @pytest.mark.asyncio + async def test_a_colliding_name_never_reaches_readline(self): + proc = _process( + _FakeShell( + [ + _frame(b"impersonating\n", channel="NOT_STDOUT"), + _frame(b'{"jsonrpc":"2.0","id":9}\n'), + ] + ) + ) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + assert ( + await asyncio.wait_for(proc.readline(), timeout=2) + == b'{"jsonrpc":"2.0","id":9}\n' + ) + + +class TestHalfCloseGuard: + """Guards PR #937: no writes to a transport whose read side is gone.""" + + @pytest.mark.asyncio + async def test_writeline_refuses_after_clean_eof(self): + """ContainerTransport.send() does not pre-check liveness. + + Without this guard an ACP request is handed to a dead socket and its + reply can never arrive, so the run stalls instead of failing. + """ + shell = _FakeShell([]) + proc = _process(shell) + await proc._drain_frames() + + with pytest.raises(RuntimeError, match="not running"): + await proc.writeline("lost-message") + + assert shell.sent == [] + + @pytest.mark.asyncio + async def test_writeline_refuses_after_reader_error(self): + shell = _FakeShell([], error=ConnectionResetError("socket died")) + proc = _process(shell) + await proc._drain_frames() + + with pytest.raises(RuntimeError, match="not running"): + await proc.writeline("lost-message") + + assert shell.sent == [] + + @pytest.mark.asyncio + async def test_writeline_works_while_the_reader_is_alive(self): + """Positive control: the guard must not break normal operation.""" + shell = _FakeShell([_frame(b"hi\n")]) + proc = _process(shell) + + await proc.writeline("ping") + + assert shell.sent == ["ping\n"] From 112d0a101ffcbe0a5fc52e9b4c94c7ec0778907a Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 13:33:51 -0700 Subject: [PATCH 08/13] fix(agentcore): close thermo review blockers --- pyproject.toml | 5 + src/benchflow/rollout/__init__.py | 7 + src/benchflow/sandbox/_base.py | 9 + src/benchflow/sandbox/agentcore.py | 340 +++++--------- src/benchflow/sandbox/agentcore_builder.py | 236 +++++++--- src/benchflow/sandbox/agentcore_image.py | 193 ++++++++ .../sandbox/agentcore_provisioning.py | 214 +++++---- src/benchflow/sandbox/process/agentcore.py | 83 +++- tests/test_agentcore_builder.py | 343 ++++++++++++++ tests/test_agentcore_parallel.py | 432 +++++------------- tests/test_agentcore_reaper.py | 297 ++++++++++++ tests/test_agentcore_sandbox.py | 51 ++- tests/test_agentcore_transport.py | 155 ++++++- tests/test_runtime_config_wired.py | 17 + uv.lock | 13 + 15 files changed, 1667 insertions(+), 728 deletions(-) create mode 100644 src/benchflow/sandbox/agentcore_image.py create mode 100644 tests/test_agentcore_reaper.py diff --git a/pyproject.toml b/pyproject.toml index 898c316a5..3c72cbdae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,9 @@ classifiers = [ [project.optional-dependencies] dev = [ "pre-commit>=3.7", + # AgentCore's CodeBuild path must reproduce Docker's ignore rules in tests + # even when the optional AWS backend is not installed. + "pathspec>=0.12", "pytest>=9.0.3", "pytest-asyncio>=0.24.0", "ruff>=0.7.0", @@ -78,6 +81,8 @@ sandbox-agentcore = [ # models the bedrock-agentcore/-control services. "bedrock-agentcore>=1.18", "boto3>=1.43.31", + # Docker-compatible build-context filtering without requiring a daemon. + "pathspec>=0.12", ] bedrock = [ "boto3>=1.40", diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 63909c2dc..7c7d1ed8d 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -956,6 +956,13 @@ async def setup(self) -> None: self._timeout = int(cfg.timeout) else: self._timeout = int(self._task.config.agent.timeout_sec or 0) + # Look on the type, not via instance getattr: permissive proxies and + # AsyncMock-based test sandboxes manufacture arbitrary attributes, + # which would turn this optional synchronous hook into an un-awaited + # fake coroutine. + configure_timeout = getattr(type(self._env), "configure_agent_timeout", None) + if callable(configure_timeout): + configure_timeout(self._env, self._timeout) _write_config( self._rollout_dir, diff --git a/src/benchflow/sandbox/_base.py b/src/benchflow/sandbox/_base.py index 80824f987..ebe58cef2 100644 --- a/src/benchflow/sandbox/_base.py +++ b/src/benchflow/sandbox/_base.py @@ -411,6 +411,15 @@ async def live_process(self, *, agent: str | None = None) -> LiveProcess: "live_process() (docker, daytona, apple-container, agentcore)." ) + def configure_agent_timeout(self, timeout_sec: int) -> None: + """Receive the rollout's effective agent wall-clock budget. + + Most backends do not need this hint. Providers whose remote session has + its own lifecycle cap can override it so the infrastructure does not + expire before BenchFlow's agent budget does. + """ + return None + # Container-level snapshot/restore (Branch substrate) # # Part of the Sandbox contract (``docs/architecture.md``): the Branch diff --git a/src/benchflow/sandbox/agentcore.py b/src/benchflow/sandbox/agentcore.py index 7c7197f8c..8be61c611 100644 --- a/src/benchflow/sandbox/agentcore.py +++ b/src/benchflow/sandbox/agentcore.py @@ -26,7 +26,7 @@ ``InvokeAgentRuntimeCommand`` against it fails with a 500 from the runtime. Serving ``GET /ping`` on port 8080 fixes it. Task images know nothing about AgentCore, so BenchFlow appends a tiny stdlib-only responder to the image - and makes it the entrypoint (see ``_PING_SHIM``). + and makes it the entrypoint (see ``agentcore_image.PING_SHIM``). 2. **Command execution and the interactive shell share one session.** A file written by ``exec()`` is visible to the agent running under ``open_shell()`` when both use the same ``runtimeSessionId``, which is what @@ -45,6 +45,7 @@ import asyncio import base64 +import binascii import io import os import shlex @@ -61,6 +62,7 @@ from benchflow.sandbox import agentcore_provisioning as provisioning from benchflow.sandbox._base import BaseSandbox, ExecResult, wrap_command_with_env_file from benchflow.sandbox._compose import compose_definition_path +from benchflow.sandbox.agentcore_image import AgentCoreImagePublisher from benchflow.sandbox.protocol import SandboxStartupError from benchflow.task.config import SandboxConfig from benchflow.task.paths import RolloutPaths @@ -76,11 +78,13 @@ _SESSION_WARMUP_ATTEMPTS = 8 _SESSION_WARMUP_BACKOFF_SEC = 4.0 _SESSION_WARMUP_MAX_BACKOFF_SEC = 30.0 +_LIFECYCLE_MIN_SEC = 60 +_LIFECYCLE_MAX_SEC = 28800 # The service caps a single command payload at 64 KB. Base64 inflates by 4/3, # and the wrapper adds shell scaffolding, so keep a wide margin. _MAX_INLINE_UPLOAD_BYTES = 24 * 1024 _MAX_DOWNLOAD_BYTES = 64 * 1024 * 1024 -# ``_PING_SHIM`` implements the plain HTTP contract (GET /ping, +# ``PING_SHIM`` implements the plain HTTP contract (GET /ping, # POST /invocations), so register the runtime as HTTP rather than relying on # whatever the service currently defaults to. _PROTOCOL_CONFIGURATION = {"serverProtocol": "HTTP"} @@ -94,47 +98,6 @@ _ENV_IDLE_TIMEOUT = "BENCHFLOW_AGENTCORE_IDLE_TIMEOUT_SEC" _ENV_MAX_LIFETIME = "BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC" -# Stdlib-only responder for the AgentCore Runtime HTTP contract. Kept -# dependency-free so it runs on any task base image that has a `python3`. -_PING_SHIM = '''\ -"""BenchFlow shim: satisfies the AgentCore Runtime HTTP contract. - -AgentCore refuses to service InvokeAgentRuntimeCommand for a session whose -container does not answer this contract, so BenchFlow injects this responder -as the image entrypoint. It does no work beyond staying alive and replying; -the agent itself is launched later over the shell WebSocket. -""" -import json -from http.server import BaseHTTPRequestHandler, HTTPServer - - -class _Handler(BaseHTTPRequestHandler): - def _reply(self, code, payload): - body = json.dumps(payload).encode() - self.send_response(code) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def do_GET(self): - if self.path.rstrip("/") in ("/ping", ""): - self._reply(200, {"status": "Healthy"}) - else: - self._reply(404, {"error": "not found"}) - - def do_POST(self): - self.rfile.read(int(self.headers.get("Content-Length") or 0)) - self._reply(200, {"result": "benchflow-sandbox"}) - - def log_message(self, *_args): - return - - -if __name__ == "__main__": - HTTPServer(("0.0.0.0", 8080), _Handler).serve_forever() -''' - class AgentCoreSandbox(BaseSandbox): """Sandbox backend for Bedrock AgentCore Runtime microVMs.""" @@ -158,16 +121,23 @@ def __init__( **kwargs, ) self.region = os.environ.get(_ENV_REGION) or _DEFAULT_REGION - self._ecr_repository = ( - os.environ.get(_ENV_ECR_REPOSITORY) or _DEFAULT_ECR_REPOSITORY - ) + ecr_repository = os.environ.get(_ENV_ECR_REPOSITORY) or _DEFAULT_ECR_REPOSITORY self._role_arn = os.environ.get(_ENV_ROLE_ARN) + self._agent_timeout_sec = 0 self.runtime_arn: str | None = None self.runtime_session_id: str | None = None self._runtime_id: str | None = None - self._account_id: str | None = None self._clients: dict[str, Any] = {} self._client_lock = threading.Lock() + self._images = AgentCoreImagePublisher( + environment_dir=environment_dir, + environment_name=environment_name, + task_env_config=task_env_config, + region=self.region, + ecr_repository=ecr_repository, + client_factory=self._client, + logger=self.logger, + ) # ---------------------------------------------------------------- config @@ -219,25 +189,27 @@ def preflight(cls) -> None: def _validate_definition(self) -> None: dockerfile = self.environment_dir / "Dockerfile" - if not dockerfile.exists() and not self.task_env_config.docker_image: + if not self.task_env_config.docker_image and ( + not os.path.lexists(dockerfile) + or dockerfile.is_symlink() + or not dockerfile.is_file() + ): raise ValueError( - f"No Dockerfile found in {self.environment_dir} and no " - "docker_image specified in task config." + f"{dockerfile} must be a regular, non-symlink Dockerfile " + "when no docker_image is specified." ) for reserved in ( provisioning.GENERATED_DOCKERFILE, provisioning.GENERATED_SHIM, ): - if (self.environment_dir / reserved).exists(): - # BenchFlow writes these into the caller's own task directory - # and removes them afterwards. A task shipping a file at either - # path would be overwritten and then permanently deleted, and - # the canonical context walk skips both names, so the collision - # would also cache-hit as if the task file did not exist. + if os.path.lexists(self.environment_dir / reserved): + # Staging owns these names and the canonical walk excludes + # them. Accepting a task file here would silently replace it + # with backend scaffolding and cache-hit as if it did not exist. raise ValueError( f"{reserved} already exists in {self.environment_dir}. " - "That path is reserved by the AgentCore backend, which " - "would overwrite and then delete it. Rename the task file." + "That path is reserved by the AgentCore backend. Rename " + "the task file." ) compose = compose_definition_path(self.environment_dir) if compose is not None: @@ -276,17 +248,9 @@ def _client(self, service: str) -> Any: self._clients[service] = cached return cached - def _resolve_account_id(self) -> str: - if self._account_id is None: - self._account_id = self._client("sts").get_caller_identity()["Account"] - return self._account_id - - def _ecr_registry(self) -> str: - return f"{self._resolve_account_id()}.dkr.ecr.{self.region}.amazonaws.com" - async def start(self, force_build: bool) -> None: try: - image_uri = await self._publish_image(force_build=force_build) + image_uri = await self._images.publish(force_build=force_build) self.runtime_arn = await self._ensure_runtime(image_uri) # AgentCore requires a session id of at least 33 characters. self.runtime_session_id = f"{uuid.uuid4()}-{uuid.uuid4().hex[:8]}" @@ -311,134 +275,6 @@ async def start(self, force_build: bool) -> None: build_timeout_sec=self.task_env_config.build_timeout_sec, ) from exc - def _image_identity(self) -> tuple[str, str]: - """``(context digest, ECR tag)`` for this task's image. - - Computed from the build context's *contents* before anything is built, - so concurrent rollouts of the same task agree on the tag without each - running a build to discover it. - """ - digest = provisioning.build_context_digest( - self.environment_dir, self._generated_dockerfile_text(), _PING_SHIM - ) - return digest, provisioning.image_tag(self.environment_name, digest) - - async def _publish_image(self, *, force_build: bool) -> str: - """Build and push the task image, returning its **immutable** URI. - - Deduplicated process-wide by content digest: a fan-out of rollouts over - one task builds and pushes exactly once, and every later run of an - unchanged task skips the build entirely on the ECR existence check. - - The value returned is ``repo@sha256:...``, not ``repo:tag``. A tag is - mutable — ``force_build``, a rebased base image, or a repository change - can all move it — so binding a runtime to a tag makes the runtime's - contents unknowable after the fact. The registry digest pins exactly - the bytes that were pushed. - """ - _digest, tag = self._image_identity() - tagged_uri = f"{self._ecr_registry()}/{self._ecr_repository}:{tag}" - cache_key = f"image:{tagged_uri}:{force_build}" - - async def _publish() -> str: - await asyncio.to_thread(self._ensure_ecr_repository) - if force_build or not await asyncio.to_thread(self._ecr_image_exists, tag): - await self._build_and_push(tagged_uri, force_build=force_build) - else: - self.logger.info("Reusing published AgentCore image %s", tagged_uri) - return await asyncio.to_thread(self._resolve_image_digest, tag) - - return await provisioning.once(cache_key, _publish) - - def _resolve_image_digest(self, tag: str) -> str: - """Resolve *tag* to the immutable ``repo@sha256:...`` reference.""" - response = self._client("ecr").describe_images( - repositoryName=self._ecr_repository, imageIds=[{"imageTag": tag}] - ) - details = response.get("imageDetails") or [] - if not details or not details[0].get("imageDigest"): - raise RuntimeError( - f"ECR did not report a digest for {self._ecr_repository}:{tag}; " - "cannot bind an AgentCore runtime to a verifiable image." - ) - return ( - f"{self._ecr_registry()}/{self._ecr_repository}@{details[0]['imageDigest']}" - ) - - async def _build_and_push(self, image_uri: str, *, force_build: bool) -> None: - builder = builders.select_builder( - self._client, - account_id=self._resolve_account_id(), - region=self.region, - ) - self.logger.info("Building %s via %s", image_uri, builder.name) - await builder.build_and_push( - builders.BuildRequest( - context_dir=self.environment_dir, - dockerfile_text=self._generated_dockerfile_text(), - shim_text=_PING_SHIM, - image_uri=image_uri, - registry=self._ecr_registry(), - region=self.region, - force_build=force_build, - timeout_sec=self.task_env_config.build_timeout_sec, - ) - ) - - def _generated_dockerfile_text(self) -> str: - """The Dockerfile BenchFlow actually builds, as text. - - Kept separate from writing it so the image digest can be computed from - the same bytes without touching the task tree. - """ - if self.task_env_config.docker_image: - base = f"FROM {self.task_env_config.docker_image}\n" - else: - base = (self.environment_dir / "Dockerfile").read_text() - return ( - base - + "\n" - + "# --- BenchFlow AgentCore runtime contract ---\n" - + "# AgentCore refuses command execution for a session whose\n" - + "# container does not answer GET /ping on :8080.\n" - + f"COPY {provisioning.GENERATED_SHIM} " - + "/opt/benchflow_agentcore_shim.py\n" - + "EXPOSE 8080\n" - + "ENTRYPOINT []\n" - + 'CMD ["python3", "/opt/benchflow_agentcore_shim.py"]\n' - ) - - def _ensure_ecr_repository(self) -> None: - from botocore.exceptions import ClientError - - ecr = self._client("ecr") - try: - ecr.create_repository(repositoryName=self._ecr_repository) - except ClientError as exc: - if exc.response["Error"]["Code"] != "RepositoryAlreadyExistsException": - raise - - def _ecr_image_exists(self, tag: str) -> bool: - """True when *tag* is present. Only a real miss counts as a miss. - - Treating access-denied or throttling as "not found" would start a - doomed build and bury the actual AWS error. - """ - from botocore.exceptions import ClientError - - try: - self._client("ecr").describe_images( - repositoryName=self._ecr_repository, imageIds=[{"imageTag": tag}] - ) - return True - except ClientError as exc: - if exc.response["Error"]["Code"] in { - "ImageNotFoundException", - "RepositoryNotFoundException", - }: - return False - raise - async def _ensure_runtime(self, image_uri: str) -> str: """Resolve the shared agent runtime for *image_uri*, creating it once. @@ -449,7 +285,7 @@ async def _ensure_runtime(self, image_uri: str) -> str: create one runtime and the first to finish deleted it while the others were still running. """ - digest, _tag = self._image_identity() + digest, _tag = self._images.identity() name = provisioning.runtime_name(self.environment_name, digest) async def _create() -> tuple[str, str]: @@ -504,6 +340,7 @@ async def _create_or_adopt_runtime( raise arn, runtime_id, current_image = existing await asyncio.to_thread(self._write_lease, control, arn) + provisioning.mark_lease_renewed(arn, time.monotonic()) if current_image != image_uri: # Adopting a runtime bound to different bytes would silently run # the wrong environment. The name encodes the build-context @@ -546,6 +383,7 @@ async def _create_or_adopt_runtime( runtime_id = created["agentRuntimeId"] await asyncio.to_thread(self._write_lease, control, created["agentRuntimeArn"]) + provisioning.mark_lease_renewed(created["agentRuntimeArn"], time.monotonic()) await asyncio.to_thread(self._wait_ready, control, runtime_id) self.logger.info( "Registered AgentCore runtime %s for %s", @@ -562,18 +400,21 @@ async def _renew_lease(self) -> None: window = float( max(lifecycle["maxLifetime"], lifecycle["idleRuntimeSessionTimeout"]) ) - now = time.monotonic() - if not provisioning.lease_needs_renewal(self.runtime_arn, window, now): - return - await asyncio.to_thread( - self._write_lease, - self._client("bedrock-agentcore-control"), - self.runtime_arn, + runtime_arn = self.runtime_arn + + async def _write() -> None: + await asyncio.to_thread( + self._write_lease, + self._client("bedrock-agentcore-control"), + runtime_arn, + ) + + await provisioning.renew_lease( + runtime_arn, + window, + _write, + monotonic=time.monotonic, ) - # Only now — a failed write must leave the throttle untouched so the - # next rollout retries instead of running on a lease that was never - # extended. - provisioning.mark_lease_renewed(self.runtime_arn, now) def _write_lease(self, control: Any, runtime_arn: str) -> None: """Mark the runtime as possibly-in-use until the session window closes. @@ -583,15 +424,16 @@ def _write_lease(self, control: Any, runtime_arn: str) -> None: the runtime's ``lastUpdatedAt`` — so control-plane age alone cannot distinguish an idle runtime from one serving a matrix right now. - The lease is the explicit contract that closes that gap: it extends to - the longest a session started now could still be alive, and cleanup - refuses to touch a runtime whose lease has not expired. Written once - per runtime per process (provisioning is single-flighted), so it costs - nothing per rollout. + The lease is the explicit contract that closes that gap. It extends + beyond the longest session window by one renewal interval: a rollout + admitted just before the throttle boundary may still live for the full + configured window, so a lease of only ``window`` seconds would expire + while that session was still legal. """ lifecycle = self._lifecycle_configuration() window = max(lifecycle["maxLifetime"], lifecycle["idleRuntimeSessionTimeout"]) - until = datetime.now(UTC) + timedelta(seconds=window) + duration = window + provisioning.lease_renewal_interval(float(window)) + until = datetime.now(UTC) + timedelta(seconds=duration) try: control.tag_resource( resourceArn=runtime_arn, @@ -676,15 +518,35 @@ def _lifecycle_configuration(self) -> dict[str, int]: as a dead transport rather than as an infrastructure error. Default the idle window to the task's own agent timeout when that is larger. """ - idle = int(os.environ.get(_ENV_IDLE_TIMEOUT) or 0) or max( - 900, int(getattr(self.task_env_config, "agent_timeout_sec", 0) or 0) - ) - lifetime = int(os.environ.get(_ENV_MAX_LIFETIME) or 0) or 28800 + + def _configured(name: str, default: int) -> int: + raw = os.environ.get(name) + if raw in {None, ""}: + value = default + else: + try: + value = int(raw) + except ValueError as exc: + raise ValueError(f"{name} must be an integer, got {raw!r}") from exc + if not _LIFECYCLE_MIN_SEC <= value <= _LIFECYCLE_MAX_SEC: + raise ValueError( + f"{name} must be between {_LIFECYCLE_MIN_SEC} and " + f"{_LIFECYCLE_MAX_SEC} seconds for AgentCore, got {value}." + ) + return value + + requested_idle = max(900, self._agent_timeout_sec) + idle = _configured(_ENV_IDLE_TIMEOUT, requested_idle) + lifetime = _configured(_ENV_MAX_LIFETIME, 28800) return { "idleRuntimeSessionTimeout": idle, "maxLifetime": lifetime, } + def configure_agent_timeout(self, timeout_sec: int) -> None: + """Keep the remote session alive for the effective rollout budget.""" + self._agent_timeout_sec = max(0, int(timeout_sec)) + def _require_role_arn(self) -> str: if self._role_arn: return self._role_arn @@ -947,7 +809,6 @@ async def upload_file(self, source_path: Path | str, target_path: str) -> None: # bounded chunks instead of one oversized command. await self._upload_via_tar( {source: PurePosixPath(target_path)}, - root=source.parent, ) return quoted = shlex.quote(target_path) @@ -976,11 +837,9 @@ async def upload_dir( } await self.exec(f"mkdir -p {shlex.quote(target_dir)}", user="root") if members: - await self._upload_via_tar(members, root=source) + await self._upload_via_tar(members) - async def _upload_via_tar( - self, members: dict[Path, PurePosixPath], root: Path - ) -> None: + async def _upload_via_tar(self, members: dict[Path, PurePosixPath]) -> None: """Ship files as a single base64 tar stream, chunked under the cap. One archive per directory keeps this O(1) commands for the common case @@ -1009,7 +868,8 @@ async def _upload_via_tar( ) result = await self.exec( - f"base64 -d {staging} | tar -xzf - -C / && rm -f {staging}", + f"set -o pipefail; trap 'rm -f {staging}' EXIT; " + f"base64 -d {staging} | tar -xzf - -C /", timeout_sec=600, user="root", ) @@ -1026,7 +886,32 @@ async def download_file(self, source_path: str, target_path: Path | str) -> None raise RuntimeError( f"AgentCore download_file failed: {(result.stderr or '')[:500]}" ) - target.write_bytes(base64.b64decode((result.stdout or "").replace("\n", ""))) + target.write_bytes( + self._decode_download_payload(result.stdout or "", kind="file") + ) + + @staticmethod + def _decode_download_payload(raw: str, *, kind: str) -> bytes: + """Decode one bounded, validated base64 command response.""" + compact = "".join(raw.split()) + max_encoded = 4 * ((_MAX_DOWNLOAD_BYTES + 2) // 3) + if len(compact) > max_encoded: + raise RuntimeError( + f"AgentCore download_{kind} encoded payload exceeds the " + f"{_MAX_DOWNLOAD_BYTES} byte cap; narrow the {kind}." + ) + try: + blob = base64.b64decode(compact, validate=True) + except (binascii.Error, ValueError) as exc: + raise RuntimeError( + f"AgentCore download_{kind} returned invalid base64" + ) from exc + if len(blob) > _MAX_DOWNLOAD_BYTES: + raise RuntimeError( + f"AgentCore download_{kind} payload {len(blob)} bytes exceeds " + f"the {_MAX_DOWNLOAD_BYTES} byte cap; narrow the {kind}." + ) + return blob async def download_dir( self, source_dir: str, target_dir: Path | str, service: str = "main" @@ -1045,11 +930,6 @@ async def download_dir( raw = (result.stdout or "").strip() if not raw: return - blob = base64.b64decode(raw) - if len(blob) > _MAX_DOWNLOAD_BYTES: - raise RuntimeError( - f"AgentCore download_dir payload {len(blob)} bytes exceeds the " - f"{_MAX_DOWNLOAD_BYTES} byte cap; narrow the directory." - ) + blob = self._decode_download_payload(raw, kind="directory") with tarfile.open(fileobj=io.BytesIO(blob), mode="r:gz") as tar: tar.extractall(target, filter="data") diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py index 727cee292..2287a93f3 100644 --- a/src/benchflow/sandbox/agentcore_builder.py +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -24,12 +24,14 @@ import json import logging import os +import shutil import subprocess +import tempfile import time import uuid import zipfile from collections.abc import Iterator -from contextlib import contextmanager +from contextlib import contextmanager, suppress from dataclasses import dataclass from pathlib import Path from typing import Any, Protocol @@ -72,22 +74,49 @@ async def build_and_push(self, request: BuildRequest) -> None: ... @contextmanager def materialized(request: BuildRequest) -> Iterator[Path]: - """Write the generated Dockerfile and shim, yield the Dockerfile, clean up. + """Yield a generated Dockerfile inside an isolated staged build context. - The build context is the task's own environment directory, so these files - must not survive the build — otherwise every run would litter the caller's - task tree. + Generated files never belong in the caller's task tree. Fixed filenames + there were vulnerable to dangling symlinks, partial-write cleanup, and + cross-process builds overwriting or unlinking each other's inputs. Each + build now receives a unique context containing exactly the canonical, + already-filtered entries plus BenchFlow's two generated files. """ - context = request.context_dir - shim = context / provisioning.GENERATED_SHIM - dockerfile = context / provisioning.GENERATED_DOCKERFILE - shim.write_text(request.shim_text) - dockerfile.write_text(request.dockerfile_text) - try: + with tempfile.TemporaryDirectory(prefix="benchflow-agentcore-context-") as raw: + context = Path(raw) + directory_modes: list[tuple[Path, int]] = [] + # Do not copy ignore-control files into the filtered staging tree: the + # local daemon/CodeBuild docker invocation would apply them a second + # time and could exclude BenchFlow's generated shim. + ignore_controls = { + ".dockerignore", + f"{provisioning.GENERATED_DOCKERFILE}.dockerignore", + } + for source, relative in provisioning.iter_context_entries(request.context_dir): + if relative in ignore_controls: + continue + target = context / relative + if source.is_dir(): + target.mkdir(parents=True, exist_ok=True) + # Apply restrictive directory modes only after descendants are + # copied; chmod(0555) here would make a valid read-only source + # directory impossible to stage. + directory_modes.append((target, source.stat().st_mode & 0o7777)) + else: + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, target, follow_symlinks=False) + + shim = context / provisioning.GENERATED_SHIM + dockerfile = context / provisioning.GENERATED_DOCKERFILE + shim.write_text(request.shim_text) + dockerfile.write_text(request.dockerfile_text) + # Generated modes are part of the image but not caller-controlled + # identity inputs, so make them deterministic across host umasks. + shim.chmod(0o644) + dockerfile.chmod(0o644) + for directory, mode in reversed(directory_modes): + directory.chmod(mode) yield dockerfile - finally: - for path in (dockerfile, shim): - path.unlink(missing_ok=True) def _run(*args: str, timeout: float | None = None, input_text: str | None = None): @@ -128,6 +157,7 @@ def __init__(self, client_factory: Any) -> None: async def build_and_push(self, request: BuildRequest) -> None: with materialized(request) as dockerfile: + context = dockerfile.parent args = [ "docker", "build", @@ -140,7 +170,7 @@ async def build_and_push(self, request: BuildRequest) -> None: ] if request.force_build: args.append("--no-cache") - args.append(str(request.context_dir)) + args.append(str(context)) result = await asyncio.to_thread(_run, *args, timeout=request.timeout_sec) if result.returncode != 0: raise RuntimeError( @@ -266,13 +296,28 @@ async def build_and_push(self, request: BuildRequest) -> None: key = f"contexts/{uuid.uuid4().hex}.zip" archive = await asyncio.to_thread(self._package, request) await asyncio.to_thread(self._ensure_bucket) - await asyncio.to_thread( - self._client("s3").put_object, - Bucket=self._bucket, - Key=key, - Body=archive, - ) try: + # The upload belongs inside the cleanup scope. S3 can commit the + # object and lose the response; deleting the key even when + # put_object raises is the only safe handling of that ambiguity. + upload = asyncio.create_task( + asyncio.to_thread( + self._client("s3").put_object, + Bucket=self._bucket, + Key=key, + Body=archive, + ) + ) + try: + await asyncio.shield(upload) + except asyncio.CancelledError: + # to_thread cancellation does not stop the underlying upload. + # Wait for it to settle before the finally block deletes the + # key, or a late successful upload can recreate the object + # after cleanup. + with suppress(Exception): + await upload + raise await asyncio.to_thread(self._ensure_project) await self._run_build(request, key) finally: @@ -303,11 +348,15 @@ def _package(self, request: BuildRequest) -> bytes: """ buffer = io.BytesIO() with ( - materialized(request), + materialized(request) as dockerfile, zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive, ): - for path, relative in provisioning.iter_context_files(request.context_dir): - archive.write(path, arcname=relative) + context = dockerfile.parent + for path, relative in provisioning.iter_context_entries(context): + archive.write( + path, + arcname=relative + "/" if path.is_dir() else relative, + ) # The generated scaffolding is excluded from the canonical walk # (it must not affect image identity), so add it explicitly — # CodeBuild only ever sees this archive. @@ -315,7 +364,7 @@ def _package(self, request: BuildRequest) -> bytes: provisioning.GENERATED_DOCKERFILE, provisioning.GENERATED_SHIM, ): - archive.write(request.context_dir / generated, arcname=generated) + archive.write(context / generated, arcname=generated) return buffer.getvalue() def _ensure_bucket(self) -> None: @@ -362,20 +411,29 @@ def _ensure_bucket(self) -> None: "RestrictPublicBuckets": True, }, ) + lifecycle_rule = { + "ID": "benchflow-expire-build-contexts", + "Status": "Enabled", + "Filter": {"Prefix": "contexts/"}, + "Expiration": {"Days": 1}, + } + try: + existing = s3.get_bucket_lifecycle_configuration(Bucket=self._bucket) + rules = [ + rule + for rule in existing.get("Rules", []) + if rule.get("ID") != lifecycle_rule["ID"] + ] + except ClientError as exc: + if exc.response["Error"]["Code"] != "NoSuchLifecycleConfiguration": + raise + rules = [] + # Preserve operator-owned lifecycle rules on a custom/shared bucket. + # put_bucket_lifecycle_configuration replaces the entire document. + rules.append(lifecycle_rule) s3.put_bucket_lifecycle_configuration( Bucket=self._bucket, - # Build contexts are consumed within minutes; expiring them keeps a - # benchmark run from accruing storage cost — and exposure — forever. - LifecycleConfiguration={ - "Rules": [ - { - "ID": "benchflow-expire-build-contexts", - "Status": "Enabled", - "Filter": {"Prefix": "contexts/"}, - "Expiration": {"Days": 1}, - } - ] - }, + LifecycleConfiguration={"Rules": rules}, ) def _ensure_project(self) -> None: @@ -400,21 +458,38 @@ def _ensure_project(self) -> None: except ClientError as exc: if exc.response["Error"]["Code"] != "ResourceAlreadyExistsException": raise + # This project is shared across runs. Reconcile its mutable + # execution contract so an old project (or an operator edit) + # cannot silently switch us to x86, disable Docker, or retain a + # stale service role. + codebuild.update_project(name=CODEBUILD_PROJECT, **config) + logger.info("Updated CodeBuild project %s", CODEBUILD_PROJECT) async def _run_build(self, request: BuildRequest, key: str) -> None: codebuild = self._client("codebuild") - started = await asyncio.to_thread( - codebuild.start_build, - projectName=CODEBUILD_PROJECT, - sourceTypeOverride="S3", - sourceLocationOverride=f"{self._bucket}/{key}", - buildspecOverride=json.dumps(_BUILDSPEC), - environmentVariablesOverride=[ - {"name": "BF_IMAGE_URI", "value": request.image_uri}, - {"name": "BF_REGISTRY", "value": request.registry}, - ], - timeoutInMinutesOverride=_CODEBUILD_TIMEOUT_MIN, + start = asyncio.create_task( + asyncio.to_thread( + codebuild.start_build, + projectName=CODEBUILD_PROJECT, + sourceTypeOverride="S3", + sourceLocationOverride=f"{self._bucket}/{key}", + buildspecOverride=json.dumps(_BUILDSPEC), + environmentVariablesOverride=[ + {"name": "BF_IMAGE_URI", "value": request.image_uri}, + {"name": "BF_REGISTRY", "value": request.registry}, + ], + timeoutInMinutesOverride=_CODEBUILD_TIMEOUT_MIN, + ) ) + cancelled_during_start: asyncio.CancelledError | None = None + try: + started = await asyncio.shield(start) + except asyncio.CancelledError as exc: + # The API may have accepted the build even though our task was + # cancelled while waiting for its response. Recover the id before + # honoring cancellation so the remote build can be stopped. + cancelled_during_start = exc + started = await start build_id = started["build"]["id"] logger.info( "CodeBuild %s building %s (arm64, no local docker)", @@ -422,22 +497,55 @@ async def _run_build(self, request: BuildRequest, key: str) -> None: request.image_uri, ) - deadline = time.monotonic() + ( - request.timeout_sec or _CODEBUILD_TIMEOUT_MIN * 60 - ) - while time.monotonic() < deadline: - await asyncio.sleep(_CODEBUILD_POLL_SEC) - builds = await asyncio.to_thread(codebuild.batch_get_builds, ids=[build_id]) - build = builds["builds"][0] - if build["buildStatus"] == "IN_PROGRESS": - continue - if build["buildStatus"] == "SUCCEEDED": - logger.info("Pushed AgentCore image %s (codebuild)", request.image_uri) - return - raise await asyncio.to_thread(self._build_failure, build, request.image_uri) - raise TimeoutError( - f"CodeBuild {build_id} did not finish within the build timeout" - ) + terminal = False + + async def _stop_unfinished() -> None: + try: + await asyncio.to_thread(codebuild.stop_build, id=build_id) + except Exception as exc: + # Preserve the timeout/cancellation/provider error that caused + # the abort, while making a potentially still-running paid + # build visible to the operator. + logger.warning( + "Could not stop unfinished CodeBuild %s: %s", build_id, exc + ) + + if cancelled_during_start is not None: + await asyncio.shield(_stop_unfinished()) + raise cancelled_during_start + + try: + deadline = time.monotonic() + ( + request.timeout_sec or _CODEBUILD_TIMEOUT_MIN * 60 + ) + while time.monotonic() < deadline: + await asyncio.sleep(_CODEBUILD_POLL_SEC) + builds = await asyncio.to_thread( + codebuild.batch_get_builds, ids=[build_id] + ) + build = builds["builds"][0] + if build["buildStatus"] == "IN_PROGRESS": + continue + terminal = True + if build["buildStatus"] == "SUCCEEDED": + logger.info( + "Pushed AgentCore image %s (codebuild)", request.image_uri + ) + return + raise await asyncio.to_thread( + self._build_failure, build, request.image_uri + ) + raise TimeoutError( + f"CodeBuild {build_id} did not finish within the build timeout" + ) + except asyncio.CancelledError: + if not terminal: + await asyncio.shield(_stop_unfinished()) + raise + except Exception: + if not terminal: + await _stop_unfinished() + raise def _build_failure(self, build: dict[str, Any], image_uri: str) -> Exception: """Turn a failed build into the most specific error we can offer.""" diff --git a/src/benchflow/sandbox/agentcore_image.py b/src/benchflow/sandbox/agentcore_image.py new file mode 100644 index 000000000..43ab08bcf --- /dev/null +++ b/src/benchflow/sandbox/agentcore_image.py @@ -0,0 +1,193 @@ +"""Image identity, build, and publication for the AgentCore backend. + +AgentCore sessions run only images already pushed to ECR. This module owns that +entire image lifecycle so the sandbox itself can focus on runtime/session and +data-plane behavior. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from benchflow.sandbox import agentcore_builder as builders +from benchflow.sandbox import agentcore_provisioning as provisioning +from benchflow.task.config import SandboxConfig + +# Stdlib-only responder for the AgentCore Runtime HTTP contract. Kept +# dependency-free so it runs on any task base image that has a ``python3``. +PING_SHIM = '''\ +"""BenchFlow shim: satisfies the AgentCore Runtime HTTP contract. + +AgentCore refuses to service InvokeAgentRuntimeCommand for a session whose +container does not answer this contract, so BenchFlow injects this responder +as the image entrypoint. It does no work beyond staying alive and replying; +the agent itself is launched later over the shell WebSocket. +""" +import json +from http.server import BaseHTTPRequestHandler, HTTPServer + + +class _Handler(BaseHTTPRequestHandler): + def _reply(self, code, payload): + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path.rstrip("/") in ("/ping", ""): + self._reply(200, {"status": "Healthy"}) + else: + self._reply(404, {"error": "not found"}) + + def do_POST(self): + self.rfile.read(int(self.headers.get("Content-Length") or 0)) + self._reply(200, {"result": "benchflow-sandbox"}) + + def log_message(self, *_args): + return + + +if __name__ == "__main__": + HTTPServer(("0.0.0.0", 8080), _Handler).serve_forever() +''' + + +class AgentCoreImagePublisher: + """Publish one task context to ECR and return its immutable image URI.""" + + def __init__( + self, + *, + environment_dir: Path, + environment_name: str, + task_env_config: SandboxConfig, + region: str, + ecr_repository: str, + client_factory: Callable[[str], Any], + logger: Any, + ) -> None: + self._environment_dir = environment_dir + self._environment_name = environment_name + self._task_env_config = task_env_config + self._region = region + self._ecr_repository = ecr_repository + self._client = client_factory + self._logger = logger + self._account_id: str | None = None + + def _resolve_account_id(self) -> str: + if self._account_id is None: + self._account_id = self._client("sts").get_caller_identity()["Account"] + return self._account_id + + def _registry(self) -> str: + return f"{self._resolve_account_id()}.dkr.ecr.{self._region}.amazonaws.com" + + def identity(self) -> tuple[str, str]: + """Return ``(context digest, ECR tag)`` for this task image.""" + digest = provisioning.build_context_digest( + self._environment_dir, + self.generated_dockerfile_text(), + PING_SHIM, + ) + return digest, provisioning.image_tag(self._environment_name, digest) + + async def publish(self, *, force_build: bool) -> str: + """Build/push the image once and return an immutable digest URI.""" + _digest, tag = self.identity() + tagged_uri = f"{self._registry()}/{self._ecr_repository}:{tag}" + cache_key = f"image:{tagged_uri}:{force_build}" + + async def _publish() -> str: + await asyncio.to_thread(self._ensure_ecr_repository) + if force_build or not await asyncio.to_thread(self._image_exists, tag): + await self._build_and_push(tagged_uri, force_build=force_build) + else: + self._logger.info("Reusing published AgentCore image %s", tagged_uri) + return await asyncio.to_thread(self._resolve_image_digest, tag) + + return await provisioning.once(cache_key, _publish) + + def _resolve_image_digest(self, tag: str) -> str: + response = self._client("ecr").describe_images( + repositoryName=self._ecr_repository, imageIds=[{"imageTag": tag}] + ) + details = response.get("imageDetails") or [] + if not details or not details[0].get("imageDigest"): + raise RuntimeError( + f"ECR did not report a digest for {self._ecr_repository}:{tag}; " + "cannot bind an AgentCore runtime to a verifiable image." + ) + return f"{self._registry()}/{self._ecr_repository}@{details[0]['imageDigest']}" + + async def _build_and_push(self, image_uri: str, *, force_build: bool) -> None: + builder = builders.select_builder( + self._client, + account_id=self._resolve_account_id(), + region=self._region, + ) + self._logger.info("Building %s via %s", image_uri, builder.name) + await builder.build_and_push( + builders.BuildRequest( + context_dir=self._environment_dir, + dockerfile_text=self.generated_dockerfile_text(), + shim_text=PING_SHIM, + image_uri=image_uri, + registry=self._registry(), + region=self._region, + force_build=force_build, + timeout_sec=self._task_env_config.build_timeout_sec, + ) + ) + + def generated_dockerfile_text(self) -> str: + """Return the Dockerfile BenchFlow builds without touching the task.""" + if self._task_env_config.docker_image: + base = f"FROM {self._task_env_config.docker_image}\n" + else: + base = provisioning.read_regular_text(self._environment_dir / "Dockerfile") + return ( + base + + "\n" + + "# --- BenchFlow AgentCore runtime contract ---\n" + + "# AgentCore refuses command execution for a session whose\n" + + "# container does not answer GET /ping on :8080.\n" + + f"COPY {provisioning.GENERATED_SHIM} " + + "/opt/benchflow_agentcore_shim.py\n" + + "EXPOSE 8080\n" + + "ENTRYPOINT []\n" + + 'CMD ["python3", "/opt/benchflow_agentcore_shim.py"]\n' + ) + + def _ensure_ecr_repository(self) -> None: + from botocore.exceptions import ClientError + + try: + self._client("ecr").create_repository(repositoryName=self._ecr_repository) + except ClientError as exc: + if exc.response["Error"]["Code"] != "RepositoryAlreadyExistsException": + raise + + def _image_exists(self, tag: str) -> bool: + """Return false only for a genuine ECR image/repository miss.""" + from botocore.exceptions import ClientError + + try: + self._client("ecr").describe_images( + repositoryName=self._ecr_repository, + imageIds=[{"imageTag": tag}], + ) + return True + except ClientError as exc: + if exc.response["Error"]["Code"] in { + "ImageNotFoundException", + "RepositoryNotFoundException", + }: + return False + raise diff --git a/src/benchflow/sandbox/agentcore_provisioning.py b/src/benchflow/sandbox/agentcore_provisioning.py index be32c5bd0..b067dbfb7 100644 --- a/src/benchflow/sandbox/agentcore_provisioning.py +++ b/src/benchflow/sandbox/agentcore_provisioning.py @@ -26,11 +26,16 @@ import asyncio import hashlib import logging +import os +import posixpath import re +import stat from collections.abc import Awaitable, Callable, Iterator from pathlib import Path from typing import Any +from pathspec import GitIgnoreSpec + logger = logging.getLogger("benchflow").getChild("agentcore") # Names generated into the build context; excluded from its digest so the @@ -59,6 +64,7 @@ _RESULTS: dict[str, Any] = {} #: runtime ARN -> monotonic time of the last lease refresh by this process. _LEASE_RENEWED: dict[str, float] = {} +_LEASE_LOCKS: dict[str, asyncio.Lock] = {} async def once[T](key: str, factory: Callable[[], Awaitable[T]]) -> T: @@ -88,6 +94,12 @@ def reset_cache() -> None: _RESULTS.clear() _KEY_LOCKS.clear() _LEASE_RENEWED.clear() + _LEASE_LOCKS.clear() + + +def lease_renewal_interval(window_seconds: float) -> float: + """How often a process may refresh one runtime's lease.""" + return max(window_seconds / 4, 60.0) def lease_needs_renewal(runtime_arn: str, window_seconds: float, now: float) -> bool: @@ -111,7 +123,7 @@ def lease_needs_renewal(runtime_arn: str, window_seconds: float, now: float) -> last = _LEASE_RENEWED.get(runtime_arn) if last is None: return True - return (now - last) >= max(window_seconds / 4, 60.0) + return (now - last) >= lease_renewal_interval(window_seconds) def mark_lease_renewed(runtime_arn: str, now: float) -> None: @@ -119,98 +131,109 @@ def mark_lease_renewed(runtime_arn: str, now: float) -> None: _LEASE_RENEWED[runtime_arn] = now -def _translate_ignore_pattern(pattern: str) -> re.Pattern[str]: - """Compile one ``.dockerignore`` pattern to a regex over relative paths. - - Follows Docker's matcher rather than approximating it: leading and trailing - separators are stripped (so ``/secret.env`` is the root file, not an - absolute path that never matches), ``*`` and ``?`` stop at a separator, - ``**`` spans zero or more path segments, and ``[...]`` character classes - are honored. Approximating any of these leaks files Docker excludes into - the CodeBuild upload, which is a credential-exposure bug because that - archive goes to S3. +async def renew_lease( + runtime_arn: str, + window_seconds: float, + write: Callable[[], Awaitable[None]], + *, + monotonic: Callable[[], float], +) -> bool: + """Single-flight a due lease refresh for one runtime. + + The due check is repeated under the per-runtime lock. Without that second + check, a fan-out of rollouts all observes the same stale timestamp before + the first AWS call completes and sends one ``TagResource`` per rollout. + The timestamp advances only after *write* succeeds, so a failed first call + leaves the immediate retry due. """ - cleaned = pattern.strip().strip("/") - segments = cleaned.split("/") - regex = "" - for index, segment in enumerate(segments): - last = index == len(segments) - 1 - if segment == "**": - regex += "(?:.*)?" if last else "(?:[^/]+/)*" - continue - regex += _translate_segment(segment) - if not last: - regex += "/" - return re.compile(f"^{regex}$") - - -def _translate_segment(segment: str) -> str: - """Translate one path segment's wildcards, including character classes.""" - out = "" - index = 0 - while index < len(segment): - char = segment[index] - if char == "*": - out += "[^/]*" - elif char == "?": - out += "[^/]" - elif char == "[": - close = segment.find("]", index + 2) - if close == -1: - # An unterminated class is a literal bracket, as in shell - # globbing — not a syntax error. - out += re.escape(char) - else: - body = segment[index + 1 : close] - negate = body[:1] in ("!", "^") - if negate: - body = body[1:] - # Keep ranges intact; escape only what would change meaning. - body = body.replace("\\", "\\\\").replace("]", "\\]") - out += f"[{'^' if negate else ''}{body}]" - index = close + 1 - continue - else: - out += re.escape(char) - index += 1 - return out + now = monotonic() + if not lease_needs_renewal(runtime_arn, window_seconds, now): + return False + async with _GLOBAL_LOCK: + lock = _LEASE_LOCKS.setdefault(runtime_arn, asyncio.Lock()) + async with lock: + now = monotonic() + if not lease_needs_renewal(runtime_arn, window_seconds, now): + return False + await write() + mark_lease_renewed(runtime_arn, monotonic()) + return True -def _dockerignore_matcher(context_dir: Path) -> Callable[[str], bool]: - """Compile ``.dockerignore`` into a predicate over context-relative paths. +def read_regular_text(path: Path) -> str: + """Read a regular file without following a final symlink. - Docker semantics: ``#`` comments, ``!`` re-includes, last matching rule - wins, and a rule matching a directory excludes everything beneath it. + Task-controlled Dockerfiles and ignore files are build inputs. Following a + symlink here would let the context escape even though the canonical walker + correctly skips symlinks. """ - ignore_file = context_dir / ".dockerignore" - if not ignore_file.is_file(): - return lambda _relative: False + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(path, flags) + except OSError as exc: + raise ValueError(f"{path} must be a regular, non-symlink file") from exc + try: + mode = os.fstat(fd).st_mode + if not stat.S_ISREG(mode): + raise ValueError(f"{path} must be a regular, non-symlink file") + with os.fdopen(fd, encoding="utf-8") as handle: + fd = -1 + return handle.read() + finally: + if fd >= 0: + os.close(fd) + + +def _active_ignore_file(context_dir: Path) -> Path | None: + """Return Docker's active ignore file for the generated Dockerfile.""" + specific = context_dir / f"{GENERATED_DOCKERFILE}.dockerignore" + default = context_dir / ".dockerignore" + for candidate in (specific, default): + if os.path.lexists(candidate): + if candidate.is_symlink() or not candidate.is_file(): + raise ValueError(f"{candidate} must be a regular, non-symlink file") + return candidate + return None + - rules: list[tuple[re.Pattern[str], bool]] = [] - for raw in ignore_file.read_text().splitlines(): +def _clean_dockerignore_lines(text: str) -> list[str]: + """Apply Docker's path-cleaning pass before wildcard matching.""" + cleaned: list[str] = [] + for raw in text.splitlines(): line = raw.strip() if not line or line.startswith("#"): continue negated = line.startswith("!") - body = line.lstrip("!").strip() - if body.strip("/"): - rules.append((_translate_ignore_pattern(body), negated)) + body = line[1:] if negated else line + # Backslash-escaped leading ``!``/``#`` are literals, not control + # markers; PathSpec understands those escapes after cleaning. + normalized = posixpath.normpath(body).strip("/") + if normalized in {"", "."}: + continue + cleaned.append(("!" if negated else "") + normalized) + return cleaned - def matches(relative: str) -> bool: - # A rule that matches an ancestor directory excludes its contents. - segments = relative.split("/") - candidates = ["/".join(segments[: i + 1]) for i in range(len(segments))] - ignored = False - for compiled, negated in rules: - if any(compiled.match(candidate) for candidate in candidates): - ignored = not negated - return ignored - return matches +def _dockerignore_matcher(context_dir: Path) -> Callable[[str], bool]: + """Compile Docker-compatible ignore rules for the generated Dockerfile. + ``GitIgnoreSpec`` supplies the mature ``**``, character-class, escaped + leading ``!``/``#``, directory, and last-match-wins semantics. Docker's + own preprocessing difference is handled above by cleaning ``.``/``..`` and + leading/trailing separators first. A Dockerfile-specific ignore file takes + precedence over the root ``.dockerignore``, just like ``docker build -f``. + """ + ignore_file = _active_ignore_file(context_dir) + if ignore_file is None: + return lambda _relative: False + spec = GitIgnoreSpec.from_lines( + _clean_dockerignore_lines(read_regular_text(ignore_file)) + ) + return spec.match_file -def iter_context_files(context_dir: Path) -> Iterator[tuple[Path, str]]: - """The canonical Docker build context: ``(absolute path, relative path)``. + +def iter_context_entries(context_dir: Path) -> Iterator[tuple[Path, str]]: + """Canonical Docker build context entries, including empty directories. Used by **both** the image digest and the CodeBuild upload so the two views cannot drift. They previously did: the local Docker daemon honored @@ -218,19 +241,34 @@ def iter_context_files(context_dir: Path) -> Iterator[tuple[Path, str]]: file, which shipped ignored files — including secrets — into S3 and also let an ignored file change the image identity. - Symlinks are skipped so a task-controlled link cannot pull host files into - the image or the upload (#411). + Symlinks and non-regular special files are skipped so a task-controlled + link cannot pull host files into the image or the upload (#411). Directory + entries are retained: an empty directory changes a real Docker context and + therefore must change both the digest and the CodeBuild ZIP. """ ignored = _dockerignore_matcher(context_dir) for path in sorted(context_dir.rglob("*")): - if path.is_symlink() or not path.is_file(): + if path.is_symlink(): continue relative = path.relative_to(context_dir).as_posix() - if relative in _GENERATED_NAMES or ignored(relative): + if relative in _GENERATED_NAMES: + continue + if path.is_dir(): + if not ignored(relative + "/"): + yield path, relative + continue + if not path.is_file() or ignored(relative): continue yield path, relative +def iter_context_files(context_dir: Path) -> Iterator[tuple[Path, str]]: + """Compatibility iterator over regular files in the canonical context.""" + for path, relative in iter_context_entries(context_dir): + if path.is_file(): + yield path, relative + + def build_context_digest( context_dir: Path, dockerfile_text: str, shim_text: str = "" ) -> str: @@ -259,13 +297,13 @@ def field(label: bytes, payload: bytes) -> None: field(b"dockerfile", dockerfile_text.encode()) field(b"shim", shim_text.encode()) - for path, relative in iter_context_files(context_dir): + for path, relative in iter_context_entries(context_dir): field(b"path", relative.encode()) - # Full permission bits, not just the executable flag: 0600 and 0644 - # produce different containers, and a secret dropped to 0644 during an - # upgrade must not reuse the old image. + field(b"kind", b"dir" if path.is_dir() else b"file") + # Full permission bits, not just the executable flag. field(b"mode", format(path.stat().st_mode & 0o7777, "04o").encode()) - field(b"blob", hashlib.sha256(path.read_bytes()).digest()) + if path.is_file(): + field(b"blob", hashlib.sha256(path.read_bytes()).digest()) return digest.hexdigest() diff --git a/src/benchflow/sandbox/process/agentcore.py b/src/benchflow/sandbox/process/agentcore.py index fcaca911e..f8f5335a8 100644 --- a/src/benchflow/sandbox/process/agentcore.py +++ b/src/benchflow/sandbox/process/agentcore.py @@ -23,7 +23,7 @@ import uuid from typing import Any -from benchflow.sandbox.process._base import LiveProcess +from benchflow.sandbox.process._base import _ENV_KEY_RE, LiveProcess logger = logging.getLogger(__name__) @@ -173,10 +173,19 @@ async def _write_env_file(self, env: dict[str, str]) -> str: text typed into the terminal, where the PTY would echo them straight back into the agent log. """ + invalid = [key for key in env if not _ENV_KEY_RE.match(key)] + if invalid: + raise ValueError( + "Invalid environment variable name(s): " + ", ".join(sorted(invalid)) + ) remote_path = f"/tmp/.benchflow_agent_env_{uuid.uuid4().hex[:16]}" body = "".join(f"export {k}={shlex.quote(v)}\n" for k, v in env.items()) result = await self._sandbox.write_text_file(remote_path, body, mode="600") if result is False: + with contextlib.suppress(Exception): + await self._sandbox.exec( + f"rm -f {shlex.quote(remote_path)}", timeout_sec=30 + ) raise RuntimeError("Failed to stage AgentCore agent env file") return remote_path @@ -188,17 +197,7 @@ async def start( ) -> None: from bedrock_agentcore.runtime import AgentCoreRuntimeClient - parts: list[str] = [] - if cwd: - parts.append(f"cd {shlex.quote(cwd)}") - if env: - remote_env_path = await self._write_env_file(env) - quoted = shlex.quote(remote_env_path) - parts.append(f". {quoted}") - parts.append(f"rm -f {quoted}") - parts.append(f"exec bash -lc {shlex.quote(command)}") - launch = " && ".join(parts) - + remote_env_path: str | None = None client = AgentCoreRuntimeClient(region=self._region) shell = client.open_shell( runtime_arn=self._runtime_arn, @@ -217,19 +216,48 @@ async def start( # ACP JSON-RPC frames routinely exceed the 4096-byte canonical-mode # line limit, and echo would feed the agent's own output back to it. marker = f"__BENCHFLOW_ACP_{uuid.uuid4().hex[:12]}__" - await self._shell.send( + await self._send( "stty raw -echo 2>/dev/null || " "stty -echo -icanon min 1 time 0 2>/dev/null || true; " f"echo '{marker}'\n" ) await self._await_marker(marker) self._clear_buffered_output() - await self._shell.send(launch + "\n") + if env: + remote_env_path = await self._write_env_file(env) + launch = self._launch_command(command, remote_env_path, cwd) + await self._send(launch + "\n") except BaseException: + if remote_env_path: + with contextlib.suppress(Exception): + await self._sandbox.exec( + f"rm -f {shlex.quote(remote_env_path)}", timeout_sec=30 + ) await self.close() raise logger.info("AgentCore shell marker seen, agent starting") + @staticmethod + def _launch_command( + command: str, remote_env_path: str | None, cwd: str | None + ) -> str: + """Build a subshell launch that cannot strand a staged env file.""" + parts: list[str] = [] + if remote_env_path: + quoted = shlex.quote(remote_env_path) + # The subshell is essential: if ``cd`` or sourcing fails, it exits + # and runs the trap. The long-lived PTY shell itself stays alive. + parts.append(f"trap 'rm -f {quoted}' EXIT") + if cwd: + parts.append(f"cd {shlex.quote(cwd)}") + if remote_env_path: + quoted = shlex.quote(remote_env_path) + parts.append(f". {quoted}") + parts.append(f"rm -f {quoted}") + parts.append("trap - EXIT") + parts.append(f"exec bash -lc {shlex.quote(command)}") + return "( " + " && ".join(parts) + " )" + async def _await_marker(self, marker: str) -> None: from benchflow.diagnostics import ( TransportClosedDiagnostic, @@ -325,16 +353,33 @@ def _closed(msg: str, diagnosis: str) -> TransportClosedError: ) return line - async def writeline(self, data: str) -> None: + def _closed_error(self, message: str): + from benchflow.diagnostics import ( + TransportClosedDiagnostic, + TransportClosedError, + ) + + return TransportClosedError( + message, + TransportClosedDiagnostic( + raw_message=message[:500], + transport_diagnosis="remote_session_killed", + ), + ) + + async def _send(self, data: str) -> None: if not self.is_running: # Half-close guard: ContainerTransport.send() does not pre-check # liveness, so without this an ACP request is handed to a socket # whose read side is already gone and the reply can never arrive. - raise RuntimeError( - "AgentCore shell is not running (reader ended); refusing to " - "send to a half-closed transport" + raise self._closed_error( + "AgentCore shell is not running (reader ended); refusing to send " + "to a half-closed transport" ) - await self._shell.send(data + "\n") + await self._shell.send(data) + + async def writeline(self, data: str) -> None: + await self._send(data + "\n") async def close(self) -> None: self._closed = True diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py index 328c0628f..5cf6a560f 100644 --- a/tests/test_agentcore_builder.py +++ b/tests/test_agentcore_builder.py @@ -9,9 +9,13 @@ from __future__ import annotations import asyncio +import multiprocessing +import os +import threading import time import zipfile from io import BytesIO +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -35,6 +39,18 @@ def _request(tmp_path, **overrides): return builders.BuildRequest(**defaults) +def _materialize_in_child(request, barrier, results): + """Hold one staged context open so another OS process overlaps it.""" + with builders.materialized(request) as dockerfile: + results.put( + ( + str(dockerfile.parent), + (dockerfile.parent / ".benchflow_agentcore_shim.py").read_text(), + ) + ) + barrier.wait(timeout=10) + + class TestBuilderSelection: def test_prefers_local_docker_when_the_daemon_is_up(self): with patch.object(builders, "docker_available", return_value=True): @@ -127,6 +143,103 @@ def test_symlinks_are_not_packaged(self, tmp_path): with zipfile.ZipFile(BytesIO(blob)) as archive: assert "link.txt" not in set(archive.namelist()) + def test_empty_directories_are_packaged_and_change_identity(self, tmp_path): + """Guards PR #937: empty directories are real Docker context entries.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + request = _request(tmp_path) + before = provisioning.build_context_digest( + tmp_path, request.dockerfile_text, request.shim_text + ) + (tmp_path / "empty").mkdir() + after = provisioning.build_context_digest( + tmp_path, request.dockerfile_text, request.shim_text + ) + + blob = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2")._package( + request + ) + with zipfile.ZipFile(BytesIO(blob)) as archive: + assert "empty/" in archive.namelist() + assert after != before + + def test_parallel_processes_use_distinct_staging_contexts(self, tmp_path): + """Guards PR #937 against cross-process scaffold clobbering.""" + if "fork" not in multiprocessing.get_all_start_methods(): + pytest.skip("requires fork to overlap local staging contexts") + context = multiprocessing.get_context("fork") + barrier = context.Barrier(2) + results = context.Queue() + first = _request(tmp_path, shim_text="first") + second = _request(tmp_path, shim_text="second") + processes = [ + context.Process( + target=_materialize_in_child, + args=(request, barrier, results), + ) + for request in (first, second) + ] + + for process in processes: + process.start() + observed = [results.get(timeout=10) for _ in processes] + for process in processes: + process.join(timeout=10) + + assert all(process.exitcode == 0 for process in processes) + assert len({path for path, _shim in observed}) == 2 + assert {shim for _path, shim in observed} == {"first", "second"} + assert all(not Path(path).exists() for path, _shim in observed) + assert not (tmp_path / ".benchflow_agentcore_shim.py").exists() + + def test_partial_scaffold_write_cleans_the_staging_context( + self, tmp_path, monkeypatch + ): + """Guards PR #937: a second generated-file write may fail.""" + request = _request(tmp_path) + original = Path.write_text + staged: list[Path] = [] + + def _write(path, data, *args, **kwargs): + if path.name == builders.provisioning.GENERATED_SHIM: + staged.append(path.parent) + if path.name == builders.provisioning.GENERATED_DOCKERFILE: + raise OSError("disk full") + return original(path, data, *args, **kwargs) + + monkeypatch.setattr(Path, "write_text", _write) + + with ( + pytest.raises(OSError, match="disk full"), + builders.materialized(request), + ): + pass + + assert staged and all(not path.exists() for path in staged) + + def test_read_only_directories_stage_with_their_original_mode(self, tmp_path): + """Guards PR #937: applying 0555 before copying children breaks staging.""" + request = _request(tmp_path) + readonly = tmp_path / "readonly" + readonly.mkdir() + (readonly / "payload.txt").write_text("payload") + os.chmod(readonly, 0o555) + + try: + with builders.materialized(request) as dockerfile: + staged = dockerfile.parent / "readonly" + assert (staged / "payload.txt").read_text() == "payload" + assert staged.stat().st_mode & 0o777 == 0o555 + finally: + os.chmod(readonly, 0o755) + + def test_generated_file_modes_do_not_depend_on_host_umask(self, tmp_path): + """Guards PR #937: generated shim mode is part of the built image.""" + with builders.materialized(_request(tmp_path)) as dockerfile: + shim = dockerfile.parent / ".benchflow_agentcore_shim.py" + assert dockerfile.stat().st_mode & 0o777 == 0o644 + assert shim.stat().st_mode & 0o777 == 0o644 + def test_buildspec_enforces_the_image_size_cap_remotely(self): """The 2 GB cap must be caught on the worker, before the push.""" commands = " ".join(builders._BUILDSPEC["phases"]["build"]["commands"]) @@ -270,6 +383,51 @@ def test_character_classes_are_honoured(self, tmp_path): # Positive control: the class must not over-match. assert "secretX.pem" in relatives + @pytest.mark.parametrize( + ("pattern", "secret_name"), + [ + ("foo/../secret.env", "secret.env"), + (r"\!secret.env", "!secret.env"), + ], + ) + def test_docker_path_cleaning_and_escapes_do_not_leak( + self, tmp_path, pattern, secret_name + ): + """Guards PR #937 with patterns verified against Docker 29.3.0.""" + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + (tmp_path / ".dockerignore").write_text(pattern + "\n") + (tmp_path / secret_name).write_text("DO_NOT_UPLOAD") + + blob = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2")._package( + _request(tmp_path) + ) + + with zipfile.ZipFile(BytesIO(blob)) as archive: + assert secret_name not in archive.namelist() + assert b"DO_NOT_UPLOAD" not in b"".join( + archive.read(name) + for name in archive.namelist() + if not name.endswith("/") + ) + + def test_generated_dockerfile_specific_ignore_takes_precedence(self, tmp_path): + """Guards PR #937: Dockerfile-specific ignore overrides the root file.""" + from benchflow.sandbox import agentcore_provisioning as provisioning + + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + (tmp_path / "secret.env").write_text("DO_NOT_UPLOAD") + (tmp_path / ".dockerignore").write_text("!secret.env\n") + (tmp_path / f"{provisioning.GENERATED_DOCKERFILE}.dockerignore").write_text( + "secret.env\n" + ) + + blob = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2")._package( + _request(tmp_path) + ) + + with zipfile.ZipFile(BytesIO(blob)) as archive: + assert "secret.env" not in archive.namelist() + def test_permission_bits_change_image_identity(self, tmp_path): """0600 and 0644 build different containers from identical bytes. @@ -336,6 +494,63 @@ def superseded_framing(dockerfile: str, shim: str) -> str: tmp_path, *first ) != provisioning.build_context_digest(tmp_path, *second) + +class TestBuildLifecycle: + """Remote and local builders fail closed and clean up after themselves.""" + + def test_bucket_hardening_preserves_existing_lifecycle_rules(self): + """Guards PR #937: hardening a shared bucket must not erase its policies.""" + s3 = MagicMock() + s3.get_bucket_lifecycle_configuration.return_value = { + "Rules": [ + { + "ID": "operator-archive", + "Status": "Enabled", + "Filter": {"Prefix": "history/"}, + "Transition": {"Days": 30, "StorageClass": "GLACIER"}, + }, + { + "ID": "benchflow-expire-build-contexts", + "Status": "Disabled", + "Filter": {"Prefix": "old/"}, + "Expiration": {"Days": 90}, + }, + ] + } + builder = builders.CodeBuildBuilder(lambda service: s3, "1", "us-west-2") + + builder._ensure_bucket() + + lifecycle = s3.put_bucket_lifecycle_configuration.call_args.kwargs[ + "LifecycleConfiguration" + ] + assert [rule["ID"] for rule in lifecycle["Rules"]] == [ + "operator-archive", + "benchflow-expire-build-contexts", + ] + assert lifecycle["Rules"][0]["Transition"]["StorageClass"] == "GLACIER" + assert lifecycle["Rules"][1]["Expiration"] == {"Days": 1} + + def test_existing_codebuild_project_is_reconciled(self, monkeypatch): + """Guards PR #937: stale shared project settings must not run a wrong build.""" + from botocore.exceptions import ClientError + + monkeypatch.setenv(builders.ENV_CODEBUILD_ROLE, "arn:aws:iam::1:role/build") + codebuild = MagicMock() + codebuild.create_project.side_effect = ClientError( + {"Error": {"Code": "ResourceAlreadyExistsException", "Message": "exists"}}, + "CreateProject", + ) + builder = builders.CodeBuildBuilder(lambda service: codebuild, "1", "us-west-2") + + builder._ensure_project() + + config = codebuild.update_project.call_args.kwargs + assert config["name"] == builders.CODEBUILD_PROJECT + assert config["environment"]["type"] == "ARM_CONTAINER" + assert config["environment"]["privilegedMode"] is True + assert config["serviceRole"] == "arn:aws:iam::1:role/build" + @pytest.mark.asyncio async def test_unmeasurable_image_is_not_pushed(self, tmp_path): """Failing open pushes an image that may exceed the hard 2 GB cap. @@ -422,3 +637,131 @@ async def _ticker(): # A blocking diagnostic starved the ticker to ~1 tick. assert ticks > 5 + + @pytest.mark.asyncio + async def test_ambiguous_upload_is_still_deleted(self, tmp_path): + """Guards PR #937 when S3 commits the upload but loses its response.""" + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + s3 = MagicMock() + s3.put_object.side_effect = TimeoutError("response lost") + + with ( + patch.object(builder, "_client", return_value=s3), + patch.object(builder, "_ensure_bucket"), + pytest.raises(TimeoutError, match="response lost"), + ): + await builder.build_and_push(_request(tmp_path)) + + s3.delete_object.assert_called_once() + + @pytest.mark.asyncio + async def test_cancelled_upload_finishes_before_context_deletion(self, tmp_path): + """Guards PR #937 against to_thread upload/delete reordering.""" + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + s3 = MagicMock() + entered = threading.Event() + release = threading.Event() + calls: list[str] = [] + + def _put(**_kwargs): + entered.set() + release.wait(timeout=10) + calls.append("put") + + s3.put_object.side_effect = _put + s3.delete_object.side_effect = lambda **_kwargs: calls.append("delete") + + with ( + patch.object(builder, "_client", return_value=s3), + patch.object(builder, "_ensure_bucket"), + ): + task = asyncio.create_task(builder.build_and_push(_request(tmp_path))) + await asyncio.to_thread(entered.wait, 10) + task.cancel() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert calls == ["put", "delete"] + + @pytest.mark.asyncio + async def test_timeout_stops_the_remote_build(self, tmp_path): + """Guards PR #937: a timed-out build must not keep running and push.""" + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + codebuild = MagicMock() + codebuild.start_build.return_value = {"build": {"id": "b:timeout"}} + request = _request(tmp_path, timeout_sec=1) + + with ( + patch.object(builder, "_client", return_value=codebuild), + patch.object( + builders, + "time", + MagicMock(monotonic=MagicMock(side_effect=[0.0, 2.0])), + ), + pytest.raises(TimeoutError, match="did not finish"), + ): + await builder._run_build(request, "contexts/x.zip") + + codebuild.stop_build.assert_called_once_with(id="b:timeout") + + @pytest.mark.asyncio + async def test_cancellation_stops_build_before_deleting_context(self, tmp_path): + """Guards PR #937: stop the consumer before deleting its S3 source.""" + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + s3 = MagicMock() + codebuild = MagicMock() + codebuild.start_build.return_value = {"build": {"id": "b:cancel"}} + codebuild.batch_get_builds.return_value = { + "builds": [{"id": "b:cancel", "buildStatus": "IN_PROGRESS"}] + } + calls: list[str] = [] + codebuild.stop_build.side_effect = lambda **_kwargs: calls.append("stop") + s3.delete_object.side_effect = lambda **_kwargs: calls.append("delete") + + def _client(service): + return codebuild if service == "codebuild" else s3 + + with ( + patch.object(builder, "_client", side_effect=_client), + patch.object(builder, "_ensure_bucket"), + patch.object(builder, "_ensure_project"), + patch.object(builders, "_CODEBUILD_POLL_SEC", 0), + ): + task = asyncio.create_task(builder.build_and_push(_request(tmp_path))) + while not codebuild.batch_get_builds.called: + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert calls == ["stop", "delete"] + + @pytest.mark.asyncio + async def test_cancellation_during_start_recovers_id_and_stops_build( + self, tmp_path + ): + """Guards PR #937 when StartBuild succeeds after local cancellation.""" + builder = builders.CodeBuildBuilder(MagicMock(), "1", "us-west-2") + codebuild = MagicMock() + entered = threading.Event() + release = threading.Event() + + def _start(**_kwargs): + entered.set() + release.wait(timeout=10) + return {"build": {"id": "b:start-race"}} + + codebuild.start_build.side_effect = _start + + with patch.object(builder, "_client", return_value=codebuild): + task = asyncio.create_task( + builder._run_build(_request(tmp_path), "contexts/x.zip") + ) + await asyncio.to_thread(entered.wait, 10) + task.cancel() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + codebuild.stop_build.assert_called_once_with(id="b:start-race") diff --git a/tests/test_agentcore_parallel.py b/tests/test_agentcore_parallel.py index 5a1ee2625..51ae4571c 100644 --- a/tests/test_agentcore_parallel.py +++ b/tests/test_agentcore_parallel.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +import time from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock, patch @@ -23,6 +24,7 @@ from benchflow.sandbox import agentcore_provisioning as provisioning from benchflow.sandbox.agentcore import AgentCoreSandbox +from benchflow.sandbox.agentcore_image import AgentCoreImagePublisher from benchflow.task.config import SandboxConfig @@ -48,7 +50,7 @@ def _sandbox(task_dir, *, name="demo-task", session="run-1"): rollout_paths=None, task_env_config=SandboxConfig(), ) - env._account_id = "123456789012" + env._images._account_id = "123456789012" return env @@ -62,13 +64,13 @@ def test_identity_is_stable_across_copies_of_the_same_task(self, tmp_path): a = _make_task(tmp_path / "first") b = _make_task(tmp_path / "second") - assert _sandbox(a)._image_identity() == _sandbox(b)._image_identity() + assert _sandbox(a)._images.identity() == _sandbox(b)._images.identity() def test_identity_changes_when_the_environment_changes(self, tmp_path): a = _make_task(tmp_path / "a") b = _make_task(tmp_path / "b", dockerfile="FROM python:3.13-slim\n") - assert _sandbox(a)._image_identity() != _sandbox(b)._image_identity() + assert _sandbox(a)._images.identity() != _sandbox(b)._images.identity() def test_identity_changes_when_a_context_file_changes(self, tmp_path): """Skills baked into the image must produce a distinct image.""" @@ -77,7 +79,7 @@ def test_identity_changes_when_a_context_file_changes(self, tmp_path): (b / "skills").mkdir() (b / "skills" / "SKILL.md").write_text("# a skill") - assert _sandbox(a)._image_identity() != _sandbox(b)._image_identity() + assert _sandbox(a)._images.identity() != _sandbox(b)._images.identity() def test_runtime_name_follows_the_image_not_the_task_name(self, tmp_path): """Guards the delete-out-from-under-you bug in the first AgentCore cut. @@ -110,18 +112,18 @@ async def _build(self, image_uri, *, force_build): await asyncio.sleep(0.01) with ( - patch.object(AgentCoreSandbox, "_ensure_ecr_repository"), - patch.object(AgentCoreSandbox, "_ecr_image_exists", return_value=False), - patch.object(AgentCoreSandbox, "_build_and_push", new=_build), + patch.object(AgentCoreImagePublisher, "_ensure_ecr_repository"), + patch.object(AgentCoreImagePublisher, "_image_exists", return_value=False), + patch.object(AgentCoreImagePublisher, "_build_and_push", new=_build), patch.object( - AgentCoreSandbox, + AgentCoreImagePublisher, "_resolve_image_digest", lambda self, tag: f"reg/repo@sha256:{tag}", ), ): sandboxes = [_sandbox(task_dir, session=f"run-{i}") for i in range(20)] uris = await asyncio.gather( - *(s._publish_image(force_build=False) for s in sandboxes) + *(s._images.publish(force_build=False) for s in sandboxes) ) assert builds["n"] == 1 @@ -159,19 +161,19 @@ async def _flaky(self, image_uri, *, force_build): raise RuntimeError("transient throttle") with ( - patch.object(AgentCoreSandbox, "_ensure_ecr_repository"), - patch.object(AgentCoreSandbox, "_ecr_image_exists", return_value=False), - patch.object(AgentCoreSandbox, "_build_and_push", new=_flaky), + patch.object(AgentCoreImagePublisher, "_ensure_ecr_repository"), + patch.object(AgentCoreImagePublisher, "_image_exists", return_value=False), + patch.object(AgentCoreImagePublisher, "_build_and_push", new=_flaky), patch.object( - AgentCoreSandbox, + AgentCoreImagePublisher, "_resolve_image_digest", lambda self, tag: f"reg/repo@sha256:{tag}", ), ): env = _sandbox(task_dir) with pytest.raises(RuntimeError, match="transient throttle"): - await env._publish_image(force_build=False) - await env._publish_image(force_build=False) + await env._images.publish(force_build=False) + await env._images.publish(force_build=False) assert calls["n"] == 2 @@ -213,293 +215,6 @@ def test_oversized_image_names_the_hard_quota(self): assert "daytona" in message -class TestReaper: - def _control(self, runtimes, tags): - control = MagicMock() - paginator = MagicMock() - paginator.paginate.return_value = [{"agentRuntimes": runtimes}] - control.get_paginator.return_value = paginator - control.list_tags_for_resource.side_effect = lambda resourceArn: { - "tags": tags.get(resourceArn, {}) - } - return control - - def test_only_benchflow_managed_runtimes_are_reaped(self): - """Never delete something another tool created in the same account.""" - from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes - - old = datetime.now(UTC) - timedelta(days=3) - runtimes = [ - {"agentRuntimeArn": "arn:mine", "agentRuntimeId": "mine", "createdAt": old}, - { - "agentRuntimeArn": "arn:other", - "agentRuntimeId": "other", - "createdAt": old, - }, - ] - expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() - control = self._control( - runtimes, - { - "arn:mine": { - provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, - provisioning.LEASE_TAG: expired, - } - }, - ) - - report = reap_stale_runtimes(control, max_age_minutes=60) - - assert report.deleted == ["mine"] - assert report.skipped_unmanaged == 1 - - def test_recent_runtimes_are_kept(self): - """A runtime from a run still in flight must survive cleanup.""" - from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes - - runtimes = [ - { - "agentRuntimeArn": "arn:mine", - "agentRuntimeId": "mine", - "createdAt": datetime.now(UTC) - timedelta(minutes=5), - } - ] - expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() - control = self._control( - runtimes, - { - "arn:mine": { - provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, - provisioning.LEASE_TAG: expired, - } - }, - ) - - report = reap_stale_runtimes(control, max_age_minutes=1440) - - assert report.deleted == [] - assert report.skipped_recent == 1 - control.delete_agent_runtime.assert_not_called() - - def test_dry_run_reports_without_deleting(self): - from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes - - runtimes = [ - { - "agentRuntimeArn": "arn:mine", - "agentRuntimeId": "mine", - "createdAt": datetime.now(UTC) - timedelta(days=3), - } - ] - expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() - control = self._control( - runtimes, - { - "arn:mine": { - provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, - provisioning.LEASE_TAG: expired, - } - }, - ) - - report = reap_stale_runtimes(control, max_age_minutes=60, dry_run=True) - - assert report.deleted == ["mine"] - control.delete_agent_runtime.assert_not_called() - - def test_unreadable_tags_fail_closed(self): - """If tags can't be read, assume it isn't ours.""" - from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes - - control = self._control( - [ - { - "agentRuntimeArn": "arn:x", - "agentRuntimeId": "x", - "createdAt": datetime.now(UTC) - timedelta(days=3), - } - ], - {}, - ) - control.list_tags_for_resource.side_effect = RuntimeError("denied") - - report = reap_stale_runtimes(control, max_age_minutes=60) - - assert report.deleted == [] - assert report.skipped_unmanaged == 1 - - -class TestCleanupCommandGating: - """``bench sandbox cleanup`` must not phone AWS unless AgentCore is set up.""" - - def test_cleanup_is_inert_without_agentcore_configuration(self, monkeypatch): - """A dev machine with AWS credentials must not trigger a live call. - - ``boto3`` ships with several unrelated extras, so importability is not - evidence that this account is being used for AgentCore runs. - """ - from benchflow.cli.sandbox import _cleanup_agentcore_runtimes - - monkeypatch.delenv("BENCHFLOW_AGENTCORE_ROLE_ARN", raising=False) - with patch("boto3.Session") as session: - assert ( - _cleanup_agentcore_runtimes(dry_run=True, max_age_minutes=60) is False - ) - session.assert_not_called() - - def test_credential_failure_degrades_instead_of_crashing(self, monkeypatch): - """Cleanup may still have Daytona work to do; don't abort the command.""" - from benchflow.cli.sandbox import _cleanup_agentcore_runtimes - - monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/x") - with patch("boto3.Session", side_effect=RuntimeError("no credentials")): - assert ( - _cleanup_agentcore_runtimes(dry_run=True, max_age_minutes=60) is False - ) - - def test_daytona_failure_does_not_block_agentcore_cleanup(self, monkeypatch): - """One backend's broken credentials must not strand the other's resources. - - The Daytona path exits when DAYTONA_API_KEY is missing; before this was - isolated it aborted the whole command, silently leaving AgentCore - runtimes to accumulate against a 100-per-account quota. - """ - import typer - - from benchflow.cli import sandbox as sandbox_cli - - calls: list[str] = [] - monkeypatch.setattr(sandbox_cli, "_daytona_sdk_available", lambda: True) - monkeypatch.setattr( - sandbox_cli, - "_cleanup_agentcore_runtimes", - lambda **kw: calls.append("agentcore") or True, - ) - fake_main = MagicMock() - fake_main._cleanup_daytona_sandboxes.side_effect = typer.Exit(1) - monkeypatch.setitem(__import__("sys").modules, "benchflow.cli.main", fake_main) - - sandbox_cli.sandbox_cleanup(dry_run=True, max_age_minutes=60) - - assert calls == ["agentcore"] - - -class TestAwsResponseShapeConformance: - """Pin the mocked shapes to the SDK's real ones. - - Every fixture in this file invents AWS responses. When an invented shape - drifts from the service model the tests keep passing while production - breaks — which is exactly how the reaper came to read a ``createdAt`` that - ``ListAgentRuntimes`` never returns, and so deleted fresh runtimes. - """ - - def _shape(self, operation, member=None): - boto3 = pytest.importorskip("boto3") - client = boto3.Session(region_name="us-west-2").client( - "bedrock-agentcore-control", - aws_access_key_id="x", - aws_secret_access_key="y", - ) - shape = client.meta.service_model.operation_model(operation).output_shape - return shape.members[member].member if member else shape - - def test_list_agent_runtimes_has_no_created_at(self): - """The field the reaper originally keyed on does not exist here.""" - members = self._shape("ListAgentRuntimes", "agentRuntimes").members - - assert "lastUpdatedAt" in members - assert "createdAt" not in members - - def test_get_agent_runtime_carries_the_bound_image(self): - """The adoption check reads this to refuse a mismatched runtime.""" - members = self._shape("GetAgentRuntime").members - - assert "agentRuntimeArtifact" in members - container = members["agentRuntimeArtifact"].members["containerConfiguration"] - assert "containerUri" in container.members - - -class TestReaperAgeHandling: - def _control(self, runtimes, tags): - control = MagicMock() - paginator = MagicMock() - paginator.paginate.return_value = [{"agentRuntimes": runtimes}] - control.get_paginator.return_value = paginator - control.list_tags_for_resource.side_effect = lambda resourceArn: { - "tags": tags.get(resourceArn, {}) - } - return control - - def _managed(self, arn, *, leased_until=None): - """Managed tags. Default lease is expired, i.e. reapable.""" - expiry = leased_until or (datetime.now(UTC) - timedelta(hours=1)) - return { - arn: { - provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, - provisioning.LEASE_TAG: expiry.isoformat(), - } - } - - def test_fresh_runtime_in_the_real_list_shape_is_kept(self): - """Reproduces the reported P1 with the shape AWS actually returns. - - ``ListAgentRuntimes`` has only ``lastUpdatedAt``; keying on - ``createdAt`` yielded None, skipped the age comparison, and selected a - minutes-old runtime for deletion under a one-day policy. - """ - from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes - - runtimes = [ - { - "agentRuntimeArn": "arn:mine", - "agentRuntimeId": "mine", - "agentRuntimeName": "bf_x", - "lastUpdatedAt": datetime.now(UTC) - timedelta(minutes=3), - "status": "READY", - } - ] - control = self._control(runtimes, self._managed("arn:mine")) - - report = reap_stale_runtimes(control, max_age_minutes=1440, dry_run=True) - - assert report.deleted == [] - assert report.skipped_recent == 1 - - def test_runtime_with_no_timestamp_is_kept(self): - """Unknown age must never be read as 'old enough to delete'.""" - from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes - - control = self._control( - [{"agentRuntimeArn": "arn:mine", "agentRuntimeId": "mine"}], - self._managed("arn:mine"), - ) - - report = reap_stale_runtimes(control, max_age_minutes=1440) - - assert report.deleted == [] - assert report.skipped_recent == 1 - control.delete_agent_runtime.assert_not_called() - - def test_genuinely_old_runtime_in_the_real_shape_is_reaped(self): - """The positive control: cleanup must still do its job.""" - from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes - - control = self._control( - [ - { - "agentRuntimeArn": "arn:mine", - "agentRuntimeId": "mine", - "lastUpdatedAt": datetime.now(UTC) - timedelta(days=3), - } - ], - self._managed("arn:mine"), - ) - - report = reap_stale_runtimes(control, max_age_minutes=1440) - - assert report.deleted == ["mine"] - - class TestRuntimeImageBinding: @pytest.mark.asyncio async def test_a_runtime_bound_to_another_image_is_updated( @@ -654,6 +369,42 @@ async def test_update_preserves_the_configured_lifecycle( } +class TestLifecycleConfiguration: + def test_effective_agent_timeout_expands_default_idle_window(self, tmp_path): + """Guards PR #937: long rollouts must not inherit the 15-minute default.""" + env = _sandbox(_make_task(tmp_path)) + env.configure_agent_timeout(1200) + + assert env._lifecycle_configuration()["idleRuntimeSessionTimeout"] == 1200 + + @pytest.mark.parametrize( + ("name", "value"), + [ + ("BENCHFLOW_AGENTCORE_IDLE_TIMEOUT_SEC", "59"), + ("BENCHFLOW_AGENTCORE_IDLE_TIMEOUT_SEC", "28801"), + ("BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC", "-1"), + ("BENCHFLOW_AGENTCORE_MAX_LIFETIME_SEC", "forever"), + ], + ) + def test_invalid_service_bounds_fail_before_an_aws_call( + self, tmp_path, monkeypatch, name, value + ): + """Guards PR #937: reject invalid AgentCore lifecycle values actionably.""" + monkeypatch.setenv(name, value) + env = _sandbox(_make_task(tmp_path)) + + with pytest.raises(ValueError, match=name): + env._lifecycle_configuration() + + def test_task_timeout_above_agentcore_limit_is_not_silently_clamped(self, tmp_path): + """Guards PR #937: a backend that cannot honor the task must say so.""" + env = _sandbox(_make_task(tmp_path)) + env.configure_agent_timeout(28801) + + with pytest.raises(ValueError, match="28800"): + env._lifecycle_configuration() + + class TestComposeRejection: def test_compose_task_is_refused_at_construction(self, tmp_path): """One container cannot host a task's side services.""" @@ -784,7 +535,7 @@ def test_negative_age_is_rejected(self): @pytest.mark.asyncio async def test_provisioning_writes_a_lease(self, tmp_path, monkeypatch): - """Without this the reaper has no activity signal at all.""" + """Guards PR #937: initial lease covers the throttle boundary.""" monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") env = _sandbox(_make_task(tmp_path)) control = MagicMock() @@ -799,7 +550,19 @@ async def test_provisioning_writes_a_lease(self, tmp_path, monkeypatch): tags = control.tag_resource.call_args.kwargs["tags"] assert provisioning.LEASE_TAG in tags - assert datetime.fromisoformat(tags[provisioning.LEASE_TAG]) > datetime.now(UTC) + remaining = ( + datetime.fromisoformat(tags[provisioning.LEASE_TAG]) - datetime.now(UTC) + ).total_seconds() + window = 28_800.0 + assert remaining == pytest.approx( + window + provisioning.lease_renewal_interval(window), abs=2 + ) + # Successful create/adopt writes mark the throttle; the first fan-out + # must not immediately send a redundant TagResource call. + assert ( + provisioning.lease_needs_renewal("arn:rt", window, time.monotonic()) + is False + ) class TestLeaseIntegrity: @@ -888,6 +651,21 @@ async def test_every_rollout_renews_when_due(self, tmp_path, monkeypatch): control.tag_resource.assert_called_once() assert provisioning.LEASE_TAG in control.tag_resource.call_args.kwargs["tags"] + @pytest.mark.asyncio + async def test_concurrent_due_rollouts_single_flight_the_lease( + self, tmp_path, monkeypatch + ): + """Guards PR #937: due-check/write/mark must be atomic per runtime.""" + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/rt") + env = _sandbox(_make_task(tmp_path)) + env.runtime_arn = "arn:rt-concurrent-renewal" + control = MagicMock() + + with patch.object(env, "_client", return_value=control): + await asyncio.gather(*(env._renew_lease() for _ in range(20))) + + control.tag_resource.assert_called_once() + class TestAdoptionContract: def _detail(self, **overrides): @@ -951,6 +729,7 @@ class TestDeprecatedCleanupAliasIsSafe: """ def test_negative_age_is_rejected_by_the_daytona_reaper(self): + """Guards PR #937 through the deprecated cleanup alias.""" from benchflow.sandbox.daytona_reaper import reap_stale_sandboxes client = MagicMock() @@ -969,6 +748,7 @@ def test_negative_age_is_rejected_by_the_daytona_reaper(self): ], ) def test_every_age_knob_rejects_negatives(self, kwargs): + """Guards PR #937: no library age parameter may accept negatives.""" from benchflow.sandbox.daytona_reaper import reap_stale_sandboxes with pytest.raises(ValueError, match="must be >= 0"): @@ -976,7 +756,7 @@ def test_every_age_knob_rejects_negatives(self, kwargs): @pytest.mark.parametrize("command", ["sandbox", "environment"]) def test_cli_rejects_negative_max_age(self, command): - """Both the current command and the deprecated alias must refuse it.""" + """Guards PR #937: both current and deprecated aliases must refuse it.""" from typer.testing import CliRunner from benchflow.cli.main import app @@ -996,7 +776,7 @@ def _sandbox_with_control(self, tmp_path, control): @pytest.mark.asyncio async def test_a_failed_renewal_retries_immediately(self, tmp_path, monkeypatch): - """The throttle recorded the attempt, not the write. + """Guards PR #937: the throttle must record the write, not the attempt. A first renewal that raised still advanced the window, so the retry that would have fixed it made zero tag_resource calls and the rollout @@ -1061,7 +841,7 @@ async def _warm(): order.append("warm") with ( - patch.object(env, "_publish_image", side_effect=_publish), + patch.object(env._images, "publish", side_effect=_publish), patch.object(env, "_ensure_runtime", side_effect=_ensure), patch.object(env, "_renew_lease", side_effect=_renew), patch.object(env, "_warm_session", side_effect=_warm), @@ -1079,11 +859,11 @@ class TestReservedPathsArePreserved: ["Dockerfile.benchflow-agentcore", ".benchflow_agentcore_shim.py"], ) def test_a_colliding_task_file_is_refused(self, tmp_path, reserved): - """These are written into the caller's tree and then deleted. + """Guards PR #937: generated names are an explicit backend contract. - A task shipping either name lost its file permanently, and because the - canonical walk skips both names the collision also cache-hit as though - the task file did not exist. + Staging no longer mutates the caller's tree, but a task shipping either + reserved name would still be excluded from the canonical context and + silently replaced by backend scaffolding without this refusal. """ task_dir = _make_task(tmp_path) original = "important task content\n" @@ -1094,3 +874,33 @@ def test_a_colliding_task_file_is_refused(self, tmp_path, reserved): # The refusal must happen before anything touches the file. assert (task_dir / reserved).read_text() == original + + @pytest.mark.parametrize( + "reserved", + ["Dockerfile.benchflow-agentcore", ".benchflow_agentcore_shim.py"], + ) + def test_a_dangling_reserved_symlink_is_refused(self, tmp_path, reserved): + """Guards PR #937: ``Path.exists`` misses dangling symlinks.""" + task_dir = _make_task(tmp_path) + external = tmp_path / "outside" / "missing" + (task_dir / reserved).symlink_to(external) + + with pytest.raises(ValueError, match="reserved"): + _sandbox(task_dir) + + assert not external.exists() + + def test_a_symlinked_dockerfile_cannot_escape_the_context(self, tmp_path): + """Guards PR #937: external Dockerfile bytes must never reach AWS.""" + task_dir = _make_task(tmp_path) + external = tmp_path / "outside.Dockerfile" + external.write_text("FROM scratch\nRUN echo EXTERNAL_SECRET\n") + dockerfile = task_dir / "Dockerfile" + dockerfile.unlink() + dockerfile.symlink_to(external) + + with pytest.raises(ValueError, match="regular, non-symlink"): + _sandbox(task_dir) + + with pytest.raises(ValueError, match="regular, non-symlink"): + provisioning.read_regular_text(dockerfile) diff --git a/tests/test_agentcore_reaper.py b/tests/test_agentcore_reaper.py new file mode 100644 index 000000000..61c32efd1 --- /dev/null +++ b/tests/test_agentcore_reaper.py @@ -0,0 +1,297 @@ +"""Cleanup and AWS response-shape regressions for AgentCore runtimes.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from unittest.mock import MagicMock, patch + +import pytest + +from benchflow.sandbox import agentcore_provisioning as provisioning + + +class TestReaper: + def _control(self, runtimes, tags): + control = MagicMock() + paginator = MagicMock() + paginator.paginate.return_value = [{"agentRuntimes": runtimes}] + control.get_paginator.return_value = paginator + control.list_tags_for_resource.side_effect = lambda resourceArn: { + "tags": tags.get(resourceArn, {}) + } + return control + + def test_only_benchflow_managed_runtimes_are_reaped(self): + """Never delete something another tool created in the same account.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + old = datetime.now(UTC) - timedelta(days=3) + runtimes = [ + {"agentRuntimeArn": "arn:mine", "agentRuntimeId": "mine", "createdAt": old}, + { + "agentRuntimeArn": "arn:other", + "agentRuntimeId": "other", + "createdAt": old, + }, + ] + expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + control = self._control( + runtimes, + { + "arn:mine": { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expired, + } + }, + ) + + report = reap_stale_runtimes(control, max_age_minutes=60) + + assert report.deleted == ["mine"] + assert report.skipped_unmanaged == 1 + + def test_recent_runtimes_are_kept(self): + """A runtime from a run still in flight must survive cleanup.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + runtimes = [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "createdAt": datetime.now(UTC) - timedelta(minutes=5), + } + ] + expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + control = self._control( + runtimes, + { + "arn:mine": { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expired, + } + }, + ) + + report = reap_stale_runtimes(control, max_age_minutes=1440) + + assert report.deleted == [] + assert report.skipped_recent == 1 + control.delete_agent_runtime.assert_not_called() + + def test_dry_run_reports_without_deleting(self): + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + runtimes = [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "createdAt": datetime.now(UTC) - timedelta(days=3), + } + ] + expired = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + control = self._control( + runtimes, + { + "arn:mine": { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expired, + } + }, + ) + + report = reap_stale_runtimes(control, max_age_minutes=60, dry_run=True) + + assert report.deleted == ["mine"] + control.delete_agent_runtime.assert_not_called() + + def test_unreadable_tags_fail_closed(self): + """If tags can't be read, assume it isn't ours.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control( + [ + { + "agentRuntimeArn": "arn:x", + "agentRuntimeId": "x", + "createdAt": datetime.now(UTC) - timedelta(days=3), + } + ], + {}, + ) + control.list_tags_for_resource.side_effect = RuntimeError("denied") + + report = reap_stale_runtimes(control, max_age_minutes=60) + + assert report.deleted == [] + assert report.skipped_unmanaged == 1 + + +class TestCleanupCommandGating: + """``bench sandbox cleanup`` must not phone AWS unless AgentCore is set up.""" + + def test_cleanup_is_inert_without_agentcore_configuration(self, monkeypatch): + """A dev machine with AWS credentials must not trigger a live call. + + ``boto3`` ships with several unrelated extras, so importability is not + evidence that this account is being used for AgentCore runs. + """ + from benchflow.cli.sandbox import _cleanup_agentcore_runtimes + + monkeypatch.delenv("BENCHFLOW_AGENTCORE_ROLE_ARN", raising=False) + with patch("boto3.Session") as session: + assert ( + _cleanup_agentcore_runtimes(dry_run=True, max_age_minutes=60) is False + ) + session.assert_not_called() + + def test_credential_failure_degrades_instead_of_crashing(self, monkeypatch): + """Cleanup may still have Daytona work to do; don't abort the command.""" + from benchflow.cli.sandbox import _cleanup_agentcore_runtimes + + monkeypatch.setenv("BENCHFLOW_AGENTCORE_ROLE_ARN", "arn:aws:iam::1:role/x") + with patch("boto3.Session", side_effect=RuntimeError("no credentials")): + assert ( + _cleanup_agentcore_runtimes(dry_run=True, max_age_minutes=60) is False + ) + + def test_daytona_failure_does_not_block_agentcore_cleanup(self, monkeypatch): + """One backend's broken credentials must not strand the other's resources. + + The Daytona path exits when DAYTONA_API_KEY is missing; before this was + isolated it aborted the whole command, silently leaving AgentCore + runtimes to accumulate against a 100-per-account quota. + """ + import typer + + from benchflow.cli import sandbox as sandbox_cli + + calls: list[str] = [] + monkeypatch.setattr(sandbox_cli, "_daytona_sdk_available", lambda: True) + monkeypatch.setattr( + sandbox_cli, + "_cleanup_agentcore_runtimes", + lambda **kw: calls.append("agentcore") or True, + ) + fake_main = MagicMock() + fake_main._cleanup_daytona_sandboxes.side_effect = typer.Exit(1) + monkeypatch.setitem(__import__("sys").modules, "benchflow.cli.main", fake_main) + + sandbox_cli.sandbox_cleanup(dry_run=True, max_age_minutes=60) + + assert calls == ["agentcore"] + + +class TestAwsResponseShapeConformance: + """Pin the mocked shapes to the SDK's real ones. + + Every fixture in this file invents AWS responses. When an invented shape + drifts from the service model the tests keep passing while production + breaks — which is exactly how the reaper came to read a ``createdAt`` that + ``ListAgentRuntimes`` never returns, and so deleted fresh runtimes. + """ + + def _shape(self, operation, member=None): + boto3 = pytest.importorskip("boto3") + client = boto3.Session(region_name="us-west-2").client( + "bedrock-agentcore-control", + aws_access_key_id="x", + aws_secret_access_key="y", + ) + shape = client.meta.service_model.operation_model(operation).output_shape + return shape.members[member].member if member else shape + + def test_list_agent_runtimes_has_no_created_at(self): + """The field the reaper originally keyed on does not exist here.""" + members = self._shape("ListAgentRuntimes", "agentRuntimes").members + + assert "lastUpdatedAt" in members + assert "createdAt" not in members + + def test_get_agent_runtime_carries_the_bound_image(self): + """The adoption check reads this to refuse a mismatched runtime.""" + members = self._shape("GetAgentRuntime").members + + assert "agentRuntimeArtifact" in members + container = members["agentRuntimeArtifact"].members["containerConfiguration"] + assert "containerUri" in container.members + + +class TestReaperAgeHandling: + def _control(self, runtimes, tags): + control = MagicMock() + paginator = MagicMock() + paginator.paginate.return_value = [{"agentRuntimes": runtimes}] + control.get_paginator.return_value = paginator + control.list_tags_for_resource.side_effect = lambda resourceArn: { + "tags": tags.get(resourceArn, {}) + } + return control + + def _managed(self, arn, *, leased_until=None): + """Managed tags. Default lease is expired, i.e. reapable.""" + expiry = leased_until or (datetime.now(UTC) - timedelta(hours=1)) + return { + arn: { + provisioning.MANAGED_TAG: provisioning.MANAGED_VALUE, + provisioning.LEASE_TAG: expiry.isoformat(), + } + } + + def test_fresh_runtime_in_the_real_list_shape_is_kept(self): + """Reproduces the reported P1 with the shape AWS actually returns. + + ``ListAgentRuntimes`` has only ``lastUpdatedAt``; keying on + ``createdAt`` yielded None, skipped the age comparison, and selected a + minutes-old runtime for deletion under a one-day policy. + """ + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + runtimes = [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "agentRuntimeName": "bf_x", + "lastUpdatedAt": datetime.now(UTC) - timedelta(minutes=3), + "status": "READY", + } + ] + control = self._control(runtimes, self._managed("arn:mine")) + + report = reap_stale_runtimes(control, max_age_minutes=1440, dry_run=True) + + assert report.deleted == [] + assert report.skipped_recent == 1 + + def test_runtime_with_no_timestamp_is_kept(self): + """Unknown age must never be read as 'old enough to delete'.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control( + [{"agentRuntimeArn": "arn:mine", "agentRuntimeId": "mine"}], + self._managed("arn:mine"), + ) + + report = reap_stale_runtimes(control, max_age_minutes=1440) + + assert report.deleted == [] + assert report.skipped_recent == 1 + control.delete_agent_runtime.assert_not_called() + + def test_genuinely_old_runtime_in_the_real_shape_is_reaped(self): + """The positive control: cleanup must still do its job.""" + from benchflow.sandbox.agentcore_reaper import reap_stale_runtimes + + control = self._control( + [ + { + "agentRuntimeArn": "arn:mine", + "agentRuntimeId": "mine", + "lastUpdatedAt": datetime.now(UTC) - timedelta(days=3), + } + ], + self._managed("arn:mine"), + ) + + report = reap_stale_runtimes(control, max_age_minutes=1440) + + assert report.deleted == ["mine"] diff --git a/tests/test_agentcore_sandbox.py b/tests/test_agentcore_sandbox.py index ead33902e..ce782be64 100644 --- a/tests/test_agentcore_sandbox.py +++ b/tests/test_agentcore_sandbox.py @@ -256,6 +256,9 @@ async def _record(command, **kwargs): # mkdir + staged chunk(s) + one extract, not 13 uploads. assert len(commands) < 13 assert any("tar -xzf" in c for c in commands) + extract = next(c for c in commands if "tar -xzf" in c) + assert "trap 'rm -f /tmp/.bf_upload_" in extract + assert "set -o pipefail" in extract @staticmethod def _extract_staged_archive(commands: list[str]) -> tarfile.TarFile: @@ -330,16 +333,41 @@ async def test_oversized_inline_write_is_refused(self, sandbox): with pytest.raises(ValueError, match="Refusing to inline"): await sandbox.write_text_file("/tmp/big", "x" * 200_000) + @pytest.mark.asyncio + async def test_oversized_file_download_is_refused_before_writing( + self, sandbox, tmp_path, monkeypatch + ): + """Guards PR #937: file downloads need the same memory cap as dirs.""" + monkeypatch.setattr("benchflow.sandbox.agentcore._MAX_DOWNLOAD_BYTES", 8) + encoded = base64.b64encode(b"012345678").decode() + + async def _exec(command, **kwargs): + return MagicMock(return_code=0, stdout=encoded, stderr="") + + target = tmp_path / "too-large.bin" + with ( + patch.object(sandbox, "exec", side_effect=_exec), + pytest.raises(RuntimeError, match="8 byte cap"), + ): + await sandbox.download_file("/remote/large", target) + + assert not target.exists() + + def test_invalid_download_base64_fails_closed(self, sandbox): + """Guards PR #937: corrupted provider output must not become a file.""" + with pytest.raises(RuntimeError, match="invalid base64"): + sandbox._decode_download_payload("not@base64", kind="file") + class TestImagePreparation: def _request(self, sandbox): from benchflow.sandbox import agentcore_builder as builders - from benchflow.sandbox.agentcore import _PING_SHIM + from benchflow.sandbox.agentcore_image import PING_SHIM return builders.BuildRequest( context_dir=sandbox.environment_dir, - dockerfile_text=sandbox._generated_dockerfile_text(), - shim_text=_PING_SHIM, + dockerfile_text=sandbox._images.generated_dockerfile_text(), + shim_text=PING_SHIM, image_uri="reg/repo:tag", registry="reg", region="us-west-2", @@ -349,7 +377,7 @@ def _request(self, sandbox): def test_generated_dockerfile_adds_the_ping_shim(self, sandbox): """Without a /ping responder the microVM 500s on every command.""" - text = sandbox._generated_dockerfile_text() + text = sandbox._images.generated_dockerfile_text() assert "FROM python:3.12-slim" in text assert "benchflow_agentcore_shim.py" in text @@ -361,15 +389,17 @@ def test_task_dockerfile_is_not_modified_in_place(self, sandbox, tmp_path): original = (tmp_path / "Dockerfile").read_text() with builders.materialized(self._request(sandbox)) as dockerfile: assert dockerfile.exists() + assert dockerfile.parent != tmp_path assert (tmp_path / "Dockerfile").read_text() == original def test_generated_files_never_survive_the_build(self, sandbox, tmp_path): - """The build context is the task's own directory; don't litter it.""" + """Guards PR #937: staging never mutates the caller's task directory.""" from benchflow.sandbox import agentcore_builder as builders before = {p.name for p in tmp_path.iterdir()} - with builders.materialized(self._request(sandbox)): - assert {p.name for p in tmp_path.iterdir()} != before + with builders.materialized(self._request(sandbox)) as dockerfile: + assert {p.name for p in tmp_path.iterdir()} == before + assert (dockerfile.parent / ".benchflow_agentcore_shim.py").is_file() assert {p.name for p in tmp_path.iterdir()} == before @@ -380,11 +410,14 @@ def test_generated_files_are_cleaned_up_after_a_failed_build( from benchflow.sandbox import agentcore_builder as builders before = {p.name for p in tmp_path.iterdir()} + staged = None with contextlib.suppress(RuntimeError): # noqa: SIM117 - with builders.materialized(self._request(sandbox)): + with builders.materialized(self._request(sandbox)) as dockerfile: + staged = dockerfile.parent raise RuntimeError("build blew up") assert {p.name for p in tmp_path.iterdir()} == before + assert staged is not None and not staged.exists() def test_docker_image_config_is_honoured(self, tmp_path): env = AgentCoreSandbox( @@ -395,7 +428,7 @@ def test_docker_image_config_is_honoured(self, tmp_path): task_env_config=SandboxConfig(docker_image="python:3.12-slim"), ) - assert "FROM python:3.12-slim" in env._generated_dockerfile_text() + assert "FROM python:3.12-slim" in env._images.generated_dockerfile_text() class TestCapabilityGating: diff --git a/tests/test_agentcore_transport.py b/tests/test_agentcore_transport.py index d4116591a..862ff10b6 100644 --- a/tests/test_agentcore_transport.py +++ b/tests/test_agentcore_transport.py @@ -10,8 +10,9 @@ from __future__ import annotations import asyncio +import re from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -50,6 +51,57 @@ async def close(self): return None +class _MarkerShell: + """Emit the nonce from the first send, then stay open or cleanly EOF.""" + + def __init__(self, *, eof_after_marker: bool, fail_launch: bool = False): + self.eof_after_marker = eof_after_marker + self.fail_launch = fail_launch + self.sent: list[str] = [] + self._marker_ready = asyncio.Event() + self._marker_emitted = False + self._closed = asyncio.Event() + self.shell_id = "fake-shell" + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._marker_emitted: + await self._marker_ready.wait() + self._marker_emitted = True + marker = re.search(r"(__BENCHFLOW_ACP_[0-9a-f]+__)", self.sent[0]) + assert marker is not None + return _frame((marker.group(1) + "\n").encode()) + if self.eof_after_marker: + raise StopAsyncIteration + await self._closed.wait() + raise StopAsyncIteration + + async def send(self, data): + if self._marker_emitted and self.fail_launch: + raise ConnectionResetError("launch socket closed") + self.sent.append(data) + self._marker_ready.set() + + async def close(self): + self._closed.set() + + +class _ShellContext: + def __init__(self, shell): + self.shell = shell + + async def __aenter__(self): + return self.shell + + +def _runtime_client(shell): + client = MagicMock() + client.open_shell.return_value = _ShellContext(shell) + return client + + def _process(shell): proc = AgentCoreProcess(MagicMock(), "arn:rt", "s" * 40, "us-west-2") proc._shell = shell @@ -186,22 +238,91 @@ async def test_secrets_are_staged_not_typed_into_the_terminal(self): assert "hunter2" in body assert proc._sandbox.write_text_file.call_args.kwargs["mode"] == "600" + @pytest.mark.asyncio + async def test_invalid_env_names_are_refused_before_staging(self): + """Guards PR #937: invalid exports fail and can strand secret files.""" + proc = _process(_FakeShell([])) + proc._sandbox.write_text_file = AsyncMock(return_value=True) + + with pytest.raises(ValueError, match=r"BAD\.KEY"): + await proc._write_env_file({"BAD.KEY": "secret"}) + + proc._sandbox.write_text_file.assert_not_called() + + @pytest.mark.asyncio + async def test_env_cleanup_is_armed_before_cwd_and_removed_before_exec(self): + """Guards PR #937: every pre-launch failure must delete staged secrets.""" + shell = _MarkerShell(eof_after_marker=False) + proc = _process(shell) + proc._sandbox.write_text_file = AsyncMock(return_value=True) + + with patch( + "bedrock_agentcore.runtime.AgentCoreRuntimeClient", + return_value=_runtime_client(shell), + ): + await proc.start("agent --serve", env={"API_KEY": "hunter2"}, cwd="/work") + + launch = shell.sent[1] + assert launch.startswith("( ") and launch.endswith(" )\n") + assert launch.index("trap 'rm -f") < launch.index("cd /work") + source = launch.index(". /tmp/.benchflow_agent_env_") + remove_after_source = launch.index("rm -f", source) + assert source < remove_after_source < launch.index("trap - EXIT") + assert "hunter2" not in launch + await proc.close() + + @pytest.mark.asyncio + async def test_failed_staging_best_effort_deletes_the_partial_file(self): + """Guards PR #937 when write succeeds but chmod reports failure.""" + proc = _process(_FakeShell([])) + proc._sandbox.write_text_file = AsyncMock(return_value=False) + proc._sandbox.exec = AsyncMock() + + with pytest.raises(RuntimeError, match="Failed to stage"): + await proc._write_env_file({"API_KEY": "hunter2"}) + + cleanup = proc._sandbox.exec.call_args.args[0] + assert cleanup.startswith("rm -f /tmp/.benchflow_agent_env_") + assert "hunter2" not in cleanup + + @pytest.mark.asyncio + async def test_failed_launch_deletes_the_staged_env_file(self): + """Guards PR #937 when the shell dies after env staging.""" + shell = _MarkerShell(eof_after_marker=False, fail_launch=True) + proc = _process(shell) + proc._sandbox.write_text_file = AsyncMock(return_value=True) + proc._sandbox.exec = AsyncMock() + + with ( + patch( + "bedrock_agentcore.runtime.AgentCoreRuntimeClient", + return_value=_runtime_client(shell), + ), + pytest.raises(ConnectionResetError, match="launch socket"), + ): + await proc.start("agent --serve", env={"API_KEY": "hunter2"}) + + cleanup = proc._sandbox.exec.call_args.args[0] + assert cleanup.startswith("rm -f /tmp/.benchflow_agent_env_") + assert "hunter2" not in cleanup + class TestAdversarialChannelNames: """Guards PR #937: typed channels are matched exactly, not by substring.""" @pytest.mark.parametrize("name", ["NOT_STDOUT", "STDOUT_METADATA", "METRICS"]) def test_colliding_typed_names_are_not_stdout(self, name): - """Substring membership admitted these; only STDOUT may be ACP input.""" + """Guards PR #937: substring collisions must not become ACP input.""" assert AgentCoreProcess._is_stdout(_frame(b"x\n", channel=name)) is False @pytest.mark.parametrize("name", ["STDOUT", "ShellChannel.STDOUT", " stdout "]) def test_genuine_stdout_is_accepted(self, name): - """Enum-qualified and case/space variants must still deliver.""" + """Guards PR #937: genuine enum/name variants must still deliver.""" assert AgentCoreProcess._is_stdout(_frame(b"x\n", channel=name)) is True @pytest.mark.asyncio async def test_a_colliding_name_never_reaches_readline(self): + """Guards PR #937 end to end against typed channel impersonation.""" proc = _process( _FakeShell( [ @@ -223,7 +344,7 @@ class TestHalfCloseGuard: @pytest.mark.asyncio async def test_writeline_refuses_after_clean_eof(self): - """ContainerTransport.send() does not pre-check liveness. + """Guards PR #937 because ContainerTransport does not pre-check liveness. Without this guard an ACP request is handed to a dead socket and its reply can never arrive, so the run stalls instead of failing. @@ -232,28 +353,48 @@ async def test_writeline_refuses_after_clean_eof(self): proc = _process(shell) await proc._drain_frames() - with pytest.raises(RuntimeError, match="not running"): + with pytest.raises(TransportClosedError, match="not running") as excinfo: await proc.writeline("lost-message") assert shell.sent == [] + assert excinfo.value.diagnostic.transport_diagnosis == "remote_session_killed" @pytest.mark.asyncio async def test_writeline_refuses_after_reader_error(self): + """Guards PR #937 when the reader ends with an exception.""" shell = _FakeShell([], error=ConnectionResetError("socket died")) proc = _process(shell) await proc._drain_frames() - with pytest.raises(RuntimeError, match="not running"): + with pytest.raises(TransportClosedError, match="not running"): await proc.writeline("lost-message") assert shell.sent == [] @pytest.mark.asyncio async def test_writeline_works_while_the_reader_is_alive(self): - """Positive control: the guard must not break normal operation.""" + """Guards PR #937: the half-close fix must preserve normal writes.""" shell = _FakeShell([_frame(b"hi\n")]) proc = _process(shell) await proc.writeline("ping") assert shell.sent == ["ping\n"] + + @pytest.mark.asyncio + async def test_marker_then_eof_refuses_the_agent_launch(self): + """Guards PR #937: startup must use the half-close checked send path.""" + shell = _MarkerShell(eof_after_marker=True) + proc = _process(shell) + + with ( + patch( + "bedrock_agentcore.runtime.AgentCoreRuntimeClient", + return_value=_runtime_client(shell), + ), + pytest.raises(TransportClosedError) as excinfo, + ): + await proc.start("agent --serve") + + assert len(shell.sent) == 1 + assert excinfo.value.diagnostic.transport_diagnosis == "remote_session_killed" diff --git a/tests/test_runtime_config_wired.py b/tests/test_runtime_config_wired.py index 4c17a1ae6..efb543c96 100644 --- a/tests/test_runtime_config_wired.py +++ b/tests/test_runtime_config_wired.py @@ -34,6 +34,10 @@ class _FakeInner: def __init__(self) -> None: self.started = 0 self.stopped = 0 + self.configured_timeout: int | None = None + + def configure_agent_timeout(self, timeout_sec: int) -> None: + self.configured_timeout = timeout_sec async def start(self, *a: Any, **kw: Any) -> None: self.started += 1 @@ -191,3 +195,16 @@ async def test_rollout_config_timeout_none_keeps_task_default( expected = int(Task(TASK_PATH).config.agent.timeout_sec or 0) assert rollout._timeout == expected + + +@pytest.mark.asyncio +async def test_effective_timeout_reaches_session_lifecycle() -> None: + """Guards PR #937: remote lifetime must follow RuntimeConfig.timeout.""" + inner = _FakeInner() + cfg = RolloutConfig(task_path=TASK_PATH, environment="agentcore", timeout=1234) + rollout = Rollout(cfg) + rollout.use_prebuilt_env(inner) + + await rollout.setup() + + assert inner.configured_timeout == 1234 diff --git a/uv.lock b/uv.lock index 98ae8e97f..259f068a3 100644 --- a/uv.lock +++ b/uv.lock @@ -353,6 +353,7 @@ deepagents = [ { name = "langchain-openai" }, ] dev = [ + { name = "pathspec" }, { name = "pre-commit" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -367,6 +368,7 @@ judge = [ sandbox-agentcore = [ { name = "bedrock-agentcore" }, { name = "boto3" }, + { name = "pathspec" }, ] sandbox-daytona = [ { name = "daytona" }, @@ -402,6 +404,8 @@ requires-dist = [ { name = "modal", marker = "extra == 'sandbox-modal'", specifier = ">=0.73" }, { name = "openai", marker = "extra == 'judge'", specifier = ">=1.40" }, { name = "packaging", specifier = ">=24" }, + { name = "pathspec", marker = "extra == 'dev'", specifier = ">=0.12" }, + { name = "pathspec", marker = "extra == 'sandbox-agentcore'", specifier = ">=0.12" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.7" }, { name = "pydantic", specifier = ">=2.7" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=9.0.3" }, @@ -2798,6 +2802,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0f/54/68a0978d1ef8502b8492099beaa6e7a0c1b32e3b5d4f677f5810cb08711c/pandas-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b2c95f8bfc1ee412bf482605d7bfd30c12d1d26bd59fdd91efeef1d4718decb1", size = 9466464, upload-time = "2026-05-11T18:54:22.754Z" }, ] +[[package]] +name = "pathspec" +version = "1.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/82/42f767fc1c1143d6fd36efb827202a2d997a375e160a71eb2888a925aac1/pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a", size = 135180, upload-time = "2026-04-27T01:46:08.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, +] + [[package]] name = "platformdirs" version = "4.9.6" From a870f31888d3404b76ec6e7d582aa0c6b7d5f081 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 13:39:31 -0700 Subject: [PATCH 09/13] fix(agentcore): filter real PTY startup noise --- src/benchflow/sandbox/process/agentcore.py | 18 ++++- tests/test_agentcore_sandbox.py | 6 +- tests/test_agentcore_transport.py | 77 ++++++++++++++++++++++ 3 files changed, 96 insertions(+), 5 deletions(-) diff --git a/src/benchflow/sandbox/process/agentcore.py b/src/benchflow/sandbox/process/agentcore.py index f8f5335a8..1a676e548 100644 --- a/src/benchflow/sandbox/process/agentcore.py +++ b/src/benchflow/sandbox/process/agentcore.py @@ -19,6 +19,7 @@ import asyncio import contextlib import logging +import re import shlex import uuid from typing import Any @@ -37,6 +38,7 @@ _START_MARKER_TIMEOUT_SEC = 180 _READLINE_TIMEOUT_ENV = "BENCHFLOW_AGENTCORE_READLINE_TIMEOUT" _READLINE_TIMEOUT_DEFAULT_SEC = 900.0 +_ANSI_CSI_RE = re.compile(rb"\x1b\[[0-?]*[ -/]*[@-~]") def _readline_timeout_sec() -> float: @@ -115,7 +117,13 @@ async def _drain_frames(self) -> None: self._partial += payload while b"\n" in self._partial: line, self._partial = self._partial.split(b"\n", 1) - await self._line_buffer.put(line.replace(b"\r", b"") + b"\n") + # Bash/readline emits bracketed-paste CSI sequences when a + # command takes over the PTY, including a standalone + # ``ESC[?2004l`` *after* the startup marker. Those bytes + # are terminal state, never ACP JSON-RPC. + line = _ANSI_CSI_RE.sub(b"", line.replace(b"\r", b"")) + if line: + await self._line_buffer.put(line + b"\n") except asyncio.CancelledError: raise except BaseException as exc: @@ -296,7 +304,13 @@ async def _await_marker(self, marker: str) -> None: transport_diagnosis="remote_session_killed", ), ) from self._failure - if marker in line.decode(errors="replace"): + # The PTY echoes the entire `stty ...; echo ''` command + # before echo suppression takes effect. Substring matching accepts + # that echoed command and releases startup too early, putting the + # prompt/command noise into ACP's JSON-RPC stream. The real marker + # is its own line, possibly wrapped in bracketed-paste ANSI codes. + text = _ANSI_CSI_RE.sub(b"", line).decode(errors="replace").strip() + if text == marker: return def _clear_buffered_output(self) -> None: diff --git a/tests/test_agentcore_sandbox.py b/tests/test_agentcore_sandbox.py index ce782be64..19c425a3d 100644 --- a/tests/test_agentcore_sandbox.py +++ b/tests/test_agentcore_sandbox.py @@ -460,9 +460,9 @@ def test_no_network_tasks_are_refused(self, tmp_path): async def test_real_agentcore_lifecycle(tmp_path): """End-to-end against the live service: build, run, transfer, tear down. - Requires AWS credentials with bedrock-agentcore and ECR access, a docker - daemon, and BENCHFLOW_AGENTCORE_ROLE_ARN pointing at the runtime execution - role. + Requires AWS credentials with bedrock-agentcore and ECR access, + BENCHFLOW_AGENTCORE_ROLE_ARN pointing at the runtime execution role, and + either Docker or BENCHFLOW_AGENTCORE_CODEBUILD_ROLE_ARN. """ (tmp_path / "Dockerfile").write_text( "FROM python:3.12-slim\nRUN echo baked > /baked.txt\n" diff --git a/tests/test_agentcore_transport.py b/tests/test_agentcore_transport.py index 862ff10b6..95c46e1db 100644 --- a/tests/test_agentcore_transport.py +++ b/tests/test_agentcore_transport.py @@ -10,14 +10,18 @@ from __future__ import annotations import asyncio +import os import re +import shlex from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest from benchflow.diagnostics import TransportClosedError +from benchflow.sandbox.agentcore import AgentCoreSandbox from benchflow.sandbox.process.agentcore import AgentCoreProcess +from benchflow.task.config import SandboxConfig def _frame(payload: bytes, channel: str = "STDOUT"): @@ -271,6 +275,35 @@ async def test_env_cleanup_is_armed_before_cwd_and_removed_before_exec(self): assert "hunter2" not in launch await proc.close() + @pytest.mark.asyncio + async def test_echoed_marker_command_does_not_complete_startup(self): + """Guards PR #937: PTY command echo is not the nonce response.""" + proc = _process(_FakeShell([])) + marker = "__BENCHFLOW_ACP_deadbeef__" + await proc._line_buffer.put( + f"root@host:/# stty raw -echo; echo '{marker}'\n".encode() + ) + await proc._line_buffer.put(f"\x1b[?2004l{marker}\r\n".encode()) + + await proc._await_marker(marker) + + assert proc._line_buffer.empty() + + @pytest.mark.asyncio + async def test_terminal_only_line_after_marker_never_reaches_acp(self): + """Guards PR #937: bracketed-paste state is not JSON-RPC input.""" + proc = _process( + _FakeShell( + [ + _frame(b"\x1b[?2004l\r\n"), + _frame(b'{"jsonrpc":"2.0","id":7}\r\n'), + ] + ) + ) + proc._reader_task = asyncio.create_task(proc._drain_frames()) + + assert await proc.readline() == b'{"jsonrpc":"2.0","id":7}\n' + @pytest.mark.asyncio async def test_failed_staging_best_effort_deletes_the_partial_file(self): """Guards PR #937 when write succeeds but chmod reports failure.""" @@ -398,3 +431,47 @@ async def test_marker_then_eof_refuses_the_agent_launch(self): assert len(shell.sent) == 1 assert excinfo.value.diagnostic.transport_diagnosis == "remote_session_killed" + + +@pytest.mark.skipif( + os.environ.get("BENCHFLOW_AGENTCORE_LIVE_TEST") != "1", + reason="live AWS test; set BENCHFLOW_AGENTCORE_LIVE_TEST=1 to run", +) +@pytest.mark.asyncio +async def test_real_agentcore_shell_transport(tmp_path): + """Guards PR #937 end to end: real WebSocket stdin/stdout/env/cwd/cleanup.""" + (tmp_path / "Dockerfile").write_text( + "FROM python:3.12-slim\nRUN echo baked > /baked.txt\n" + ) + env = AgentCoreSandbox( + environment_dir=tmp_path, + environment_name="live-canary", + session_id="live-transport", + rollout_paths=None, + task_env_config=SandboxConfig(), + ) + await env.start(force_build=False) + process = await env.live_process() + program = ( + "import os,sys;" + "print('READY:'+os.getcwd()+':'+os.environ['BF_LIVE_SECRET'], flush=True);" + "[(print('ECHO:'+line.rstrip('\\\\n'), flush=True)) for line in sys.stdin]" + ) + try: + await process.start( + f"python3 -u -c {shlex.quote(program)}", + env={"BF_LIVE_SECRET": "value with spaces"}, + cwd="/tmp", + ) + assert await process.readline() == b"READY:/tmp:value with spaces\n" + await process.writeline("hello-agentcore") + assert await process.readline() == b"ECHO:hello-agentcore\n" + + staged = await env.exec( + "find /tmp -maxdepth 1 -name '.benchflow_agent_env_*' -print" + ) + assert staged.return_code == 0 + assert not (staged.stdout or "").strip() + finally: + await process.close() + await env.stop(delete=True) From ad43e7f491423f6c337bf1f21e92dc3b059b8b84 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 13:50:21 -0700 Subject: [PATCH 10/13] fix(agentcore): retry concurrent S3 hardening --- src/benchflow/sandbox/agentcore_builder.py | 28 ++++++++++++++++++++-- tests/test_agentcore_builder.py | 25 +++++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py index 2287a93f3..de25e6ffb 100644 --- a/src/benchflow/sandbox/agentcore_builder.py +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -50,6 +50,7 @@ CODEBUILD_COMPUTE = "BUILD_GENERAL1_LARGE" _CODEBUILD_POLL_SEC = 10 _CODEBUILD_TIMEOUT_MIN = 60 +_S3_CONTROL_RETRY_DELAYS_SEC = (0.25, 0.5, 1.0, 2.0) @dataclass(frozen=True) @@ -402,7 +403,8 @@ def _ensure_bucket(self) -> None: }: raise - s3.put_public_access_block( + self._retry_s3_mutation( + s3.put_public_access_block, Bucket=self._bucket, PublicAccessBlockConfiguration={ "BlockPublicAcls": True, @@ -431,11 +433,33 @@ def _ensure_bucket(self) -> None: # Preserve operator-owned lifecycle rules on a custom/shared bucket. # put_bucket_lifecycle_configuration replaces the entire document. rules.append(lifecycle_rule) - s3.put_bucket_lifecycle_configuration( + self._retry_s3_mutation( + s3.put_bucket_lifecycle_configuration, Bucket=self._bucket, LifecycleConfiguration={"Rules": rules}, ) + @staticmethod + def _retry_s3_mutation(operation: Any, **kwargs: Any) -> None: + """Retry S3's transient control-plane conflict across build processes.""" + from botocore.exceptions import ClientError + + for attempt, delay in enumerate((*_S3_CONTROL_RETRY_DELAYS_SEC, None)): + try: + operation(**kwargs) + return + except ClientError as exc: + code = exc.response["Error"]["Code"] + if code != "OperationAborted" or delay is None: + raise + logger.warning( + "S3 control-plane mutation conflicted (attempt %d); " + "retrying in %.2fs", + attempt + 1, + delay, + ) + time.sleep(delay) + def _ensure_project(self) -> None: from botocore.exceptions import ClientError diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py index 5cf6a560f..d4305a939 100644 --- a/tests/test_agentcore_builder.py +++ b/tests/test_agentcore_builder.py @@ -498,6 +498,31 @@ def superseded_framing(dockerfile: str, shim: str) -> str: class TestBuildLifecycle: """Remote and local builders fail closed and clean up after themselves.""" + def test_s3_control_conflict_is_retried(self): + """Guards PR #937: parallel builders hit live AWS OperationAborted.""" + from botocore.exceptions import ClientError + + operation = MagicMock( + side_effect=[ + ClientError( + { + "Error": { + "Code": "OperationAborted", + "Message": "conflicting conditional operation", + } + }, + "PutBucketLifecycleConfiguration", + ), + None, + ] + ) + + with patch.object(builders.time, "sleep") as sleep: + builders.CodeBuildBuilder._retry_s3_mutation(operation, Bucket="build") + + assert operation.call_count == 2 + sleep.assert_called_once_with(builders._S3_CONTROL_RETRY_DELAYS_SEC[0]) + def test_bucket_hardening_preserves_existing_lifecycle_rules(self): """Guards PR #937: hardening a shared bucket must not erase its policies.""" s3 = MagicMock() From 4cebc12ce7080621fa27eb336627d27d1f5ebf03 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 13:51:40 -0700 Subject: [PATCH 11/13] fix(agentcore): jitter S3 conflict retries --- src/benchflow/sandbox/agentcore_builder.py | 9 ++++++--- tests/test_agentcore_builder.py | 9 ++++++++- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/benchflow/sandbox/agentcore_builder.py b/src/benchflow/sandbox/agentcore_builder.py index de25e6ffb..2c729944b 100644 --- a/src/benchflow/sandbox/agentcore_builder.py +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -24,6 +24,7 @@ import json import logging import os +import random import shutil import subprocess import tempfile @@ -50,7 +51,8 @@ CODEBUILD_COMPUTE = "BUILD_GENERAL1_LARGE" _CODEBUILD_POLL_SEC = 10 _CODEBUILD_TIMEOUT_MIN = 60 -_S3_CONTROL_RETRY_DELAYS_SEC = (0.25, 0.5, 1.0, 2.0) +_S3_CONTROL_RETRY_DELAYS_SEC = (0.25, 0.5, 1.0, 2.0, 4.0, 8.0) +_S3_RETRY_JITTER = random.SystemRandom() @dataclass(frozen=True) @@ -452,13 +454,14 @@ def _retry_s3_mutation(operation: Any, **kwargs: Any) -> None: code = exc.response["Error"]["Code"] if code != "OperationAborted" or delay is None: raise + jittered = _S3_RETRY_JITTER.uniform(delay * 0.5, delay * 1.5) logger.warning( "S3 control-plane mutation conflicted (attempt %d); " "retrying in %.2fs", attempt + 1, - delay, + jittered, ) - time.sleep(delay) + time.sleep(jittered) def _ensure_project(self) -> None: from botocore.exceptions import ClientError diff --git a/tests/test_agentcore_builder.py b/tests/test_agentcore_builder.py index d4305a939..5f0c2af82 100644 --- a/tests/test_agentcore_builder.py +++ b/tests/test_agentcore_builder.py @@ -517,7 +517,14 @@ def test_s3_control_conflict_is_retried(self): ] ) - with patch.object(builders.time, "sleep") as sleep: + with ( + patch.object( + builders._S3_RETRY_JITTER, + "uniform", + return_value=builders._S3_CONTROL_RETRY_DELAYS_SEC[0], + ), + patch.object(builders.time, "sleep") as sleep, + ): builders.CodeBuildBuilder._retry_s3_mutation(operation, Bucket="build") assert operation.call_count == 2 From 2f4de9f9c9b6dbe702b6b4a04124efbf7fe2f860 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 14:02:25 -0700 Subject: [PATCH 12/13] test(agentcore): stub optional SDK in default suite --- tests/test_agentcore_transport.py | 33 +++++++++++++++++++------------ 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/tests/test_agentcore_transport.py b/tests/test_agentcore_transport.py index 95c46e1db..67504fd2b 100644 --- a/tests/test_agentcore_transport.py +++ b/tests/test_agentcore_transport.py @@ -13,7 +13,8 @@ import os import re import shlex -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -106,6 +107,21 @@ def _runtime_client(shell): return client +def _runtime_sdk(shell): + """Install the optional AgentCore SDK modules for one transport test.""" + package = ModuleType("bedrock_agentcore") + runtime = ModuleType("bedrock_agentcore.runtime") + runtime.AgentCoreRuntimeClient = MagicMock(return_value=_runtime_client(shell)) + package.runtime = runtime + return patch.dict( + sys.modules, + { + "bedrock_agentcore": package, + "bedrock_agentcore.runtime": runtime, + }, + ) + + def _process(shell): proc = AgentCoreProcess(MagicMock(), "arn:rt", "s" * 40, "us-west-2") proc._shell = shell @@ -260,10 +276,7 @@ async def test_env_cleanup_is_armed_before_cwd_and_removed_before_exec(self): proc = _process(shell) proc._sandbox.write_text_file = AsyncMock(return_value=True) - with patch( - "bedrock_agentcore.runtime.AgentCoreRuntimeClient", - return_value=_runtime_client(shell), - ): + with _runtime_sdk(shell): await proc.start("agent --serve", env={"API_KEY": "hunter2"}, cwd="/work") launch = shell.sent[1] @@ -327,10 +340,7 @@ async def test_failed_launch_deletes_the_staged_env_file(self): proc._sandbox.exec = AsyncMock() with ( - patch( - "bedrock_agentcore.runtime.AgentCoreRuntimeClient", - return_value=_runtime_client(shell), - ), + _runtime_sdk(shell), pytest.raises(ConnectionResetError, match="launch socket"), ): await proc.start("agent --serve", env={"API_KEY": "hunter2"}) @@ -421,10 +431,7 @@ async def test_marker_then_eof_refuses_the_agent_launch(self): proc = _process(shell) with ( - patch( - "bedrock_agentcore.runtime.AgentCoreRuntimeClient", - return_value=_runtime_client(shell), - ), + _runtime_sdk(shell), pytest.raises(TransportClosedError) as excinfo, ): await proc.start("agent --serve") From 3d91c71619e03aedfd4ffb9ce026099a21f8c70b Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 28 Jul 2026 14:05:20 -0700 Subject: [PATCH 13/13] docs(cli): describe multi-backend sandbox cleanup --- src/benchflow/cli/sandbox.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/benchflow/cli/sandbox.py b/src/benchflow/cli/sandbox.py index 570cb4353..1f825d6f9 100644 --- a/src/benchflow/cli/sandbox.py +++ b/src/benchflow/cli/sandbox.py @@ -239,5 +239,5 @@ def sandbox_cleanup_cmd( ), ] = 1440, ) -> None: - """Clean up orphaned Daytona sandboxes.""" + """Clean up orphaned sandboxes and stale shared runtimes.""" sandbox_cleanup(dry_run=dry_run, max_age_minutes=max_age_minutes)