diff --git a/docs/running-benchmarks.md b/docs/running-benchmarks.md index 433b5469c..8816ff13e 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,104 @@ 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. + +#### 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 +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. + +**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. + +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 | +|---|---|---| +| 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/pyproject.toml b/pyproject.toml index c6a6fb256..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", @@ -70,6 +73,17 @@ 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", + # Docker-compatible build-context filtering without requiring a daemon. + "pathspec>=0.12", +] 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/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/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/cli/sandbox.py b/src/benchflow/cli/sandbox.py index 20fb3dd44..1f825d6f9 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: @@ -173,8 +231,13 @@ 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.""" + """Clean up orphaned sandboxes and stale shared runtimes.""" sandbox_cleanup(dry_run=dry_run, max_age_minutes=max_age_minutes) 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 aa627013b..ebe58cef2 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,37 @@ 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)." + ) + + 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/_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 new file mode 100644 index 000000000..8be61c611 --- /dev/null +++ b/src/benchflow/sandbox/agentcore.py @@ -0,0 +1,935 @@ +"""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. + +**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: + +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 ``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 + 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 binascii +import io +import os +import shlex +import tarfile +import threading +import time +import uuid +from datetime import UTC, datetime, timedelta +from pathlib import Path, PurePosixPath +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._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 + +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 +_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, +# 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" +_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" + + +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 + 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._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 + + @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 + + # 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 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: + dockerfile = self.environment_dir / "Dockerfile" + 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"{dockerfile} must be a regular, non-symlink Dockerfile " + "when no docker_image is specified." + ) + for reserved in ( + provisioning.GENERATED_DOCKERFILE, + provisioning.GENERATED_SHIM, + ): + 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. 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 + # 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 + + 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 + + async def start(self, force_build: bool) -> None: + try: + 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]}" + # 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)", + 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 _ensure_runtime(self, image_uri: str) -> str: + """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._images.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) + + # 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 + + 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") + try: + created = await asyncio.to_thread( + control.create_agent_runtime, + agentRuntimeName=name, + agentRuntimeArtifact={ + "containerConfiguration": {"containerUri": image_uri} + }, + roleArn=self._require_role_arn(), + networkConfiguration=_NETWORK_CONFIGURATION, + 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, 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 + # 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=_NETWORK_CONFIGURATION, + 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_adopted_runtime, + control, + 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 + + 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", + created["agentRuntimeArn"], + image_uri, + ) + 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"]) + ) + 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, + ) + + 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 + 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"]) + duration = window + provisioning.lease_renewal_interval(float(window)) + until = datetime.now(UTC) + timedelta(seconds=duration) + try: + control.tag_resource( + resourceArn=runtime_arn, + tags={provisioning.LEASE_TAG: until.isoformat()}, + ) + except Exception as 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( + control: Any, + 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. + + 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 {} + 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, + ) + actual = detail.get("lifecycleConfiguration") or {} + 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} does not match this run's " + f"configuration (found, expected): {drifted}. Refusing to adopt " + "a runtime with a different contract.", + sandbox_id=runtime_id, + ) + + 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. + """ + + 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 + 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 _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: + """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. + + 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( + 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 + + # ------------------------------------------------------------------ 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)}, + ) + 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) + + 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 + 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"set -o pipefail; trap 'rm -f {staging}' EXIT; " + f"base64 -d {staging} | tar -xzf - -C /", + 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( + 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" + ) -> 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 = 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 new file mode 100644 index 000000000..2c729944b --- /dev/null +++ b/src/benchflow/sandbox/agentcore_builder.py @@ -0,0 +1,643 @@ +"""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 random +import shutil +import subprocess +import tempfile +import time +import uuid +import zipfile +from collections.abc import Iterator +from contextlib import contextmanager, suppress +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 +_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) +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]: + """Yield a generated Dockerfile inside an isolated staged build context. + + 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. + """ + 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 + + +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: + context = dockerfile.parent + args = [ + "docker", + "build", + "--platform", + "linux/arm64", + "-f", + str(dockerfile), + "-t", + request.image_uri, + ] + if request.force_build: + args.append("--no-cache") + args.append(str(context)) + 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 + ) + from benchflow.sandbox.protocol import SandboxStartupError + + raw = (inspect.stdout or "").strip() + if inspect.returncode != 0 or not raw.isdigit(): + # 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: + 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)', + '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", + ] + }, + }, +} + + +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) + 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: + try: + await asyncio.to_thread( + self._client("s3").delete_object, Bucket=self._bucket, Key=key + ) + except Exception as 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. + + 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) as dockerfile, + zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive, + ): + 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. + for generated in ( + provisioning.GENERATED_DOCKERFILE, + provisioning.GENERATED_SHIM, + ): + archive.write(context / generated, arcname=generated) + 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) + except ClientError as exc: + 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 + 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 create_exc: + if create_exc.response["Error"]["Code"] not in { + "BucketAlreadyOwnedByYou", + "BucketAlreadyExists", + }: + raise + + self._retry_s3_mutation( + s3.put_public_access_block, + Bucket=self._bucket, + PublicAccessBlockConfiguration={ + "BlockPublicAcls": True, + "IgnorePublicAcls": True, + "BlockPublicPolicy": True, + "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) + 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 + 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, + jittered, + ) + time.sleep(jittered) + + 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 + # 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") + 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)", + build_id, + request.image_uri, + ) + + 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.""" + 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/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 new file mode 100644 index 000000000..b067dbfb7 --- /dev/null +++ b/src/benchflow/sandbox/agentcore_provisioning.py @@ -0,0 +1,362 @@ +"""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 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 +# 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" +#: 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. +MAX_IMAGE_MB = 2048 + +_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] = {} +_LEASE_LOCKS: dict[str, asyncio.Lock] = {} + + +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() + _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: + """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). + + 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. + """ + last = _LEASE_RENEWED.get(runtime_arn) + if last is None: + return True + return (now - last) >= lease_renewal_interval(window_seconds) + + +def mark_lease_renewed(runtime_arn: str, now: float) -> None: + """Record a lease write that succeeded.""" + _LEASE_RENEWED[runtime_arn] = now + + +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. + """ + 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 read_regular_text(path: Path) -> str: + """Read a regular file without following a final symlink. + + 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. + """ + 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 + + +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[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 _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_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 + ``.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 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(): + continue + relative = path.relative_to(context_dir).as_posix() + 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: + """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() + + 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_entries(context_dir): + field(b"path", relative.encode()) + 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()) + if path.is_file(): + field(b"blob", 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..7657b13c2 --- /dev/null +++ b/src/benchflow/sandbox/agentcore_reaper.py @@ -0,0 +1,173 @@ +"""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 collections.abc import Mapping +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from typing import Any + +from benchflow.sandbox.agentcore_provisioning import ( + LEASE_TAG, + 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 + 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"active={self.skipped_active} errors={len(self.errors)}" + ) + + +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 _read_tags(control: Any, arn: str) -> dict[str, str] | None: + """Tags for *arn*, or None when they cannot be read. + + None means "cannot prove anything about this runtime", which every caller + must treat as a reason to leave it alone. + """ + try: + 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: + # 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: + 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( + 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*. + + 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() + 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", []): + report.scanned += 1 + arn = runtime.get("agentRuntimeArn") + runtime_id = runtime.get("agentRuntimeId") + if not arn or not runtime_id: + continue + 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: + # 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) + 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/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/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/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..1a676e548 --- /dev/null +++ b/src/benchflow/sandbox/process/agentcore.py @@ -0,0 +1,419 @@ +"""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 re +import shlex +import uuid +from typing import Any + +from benchflow.sandbox.process._base import _ENV_KEY_RE, LiveProcess + +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 +_ANSI_CSI_RE = re.compile(rb"\x1b\[[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[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: + """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 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() + if not payload: + continue + self._partial += payload + while b"\n" in self._partial: + line, self._partial = self._partial.split(b"\n", 1) + # 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: + 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. + + 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) + 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 + # 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 + 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. + + 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. + """ + 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 + + async def start( + self, + command: str, + env: dict[str, str] | None = None, + cwd: str | None = None, + ) -> None: + from bedrock_agentcore.runtime import AgentCoreRuntimeClient + + remote_env_path: str | None = None + 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._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() + 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, + 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 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 + # 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: + """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): + 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 ( + 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: + line = await asyncio.wait_for(self._line_buffer.get(), timeout=timeout) + except TimeoutError as e: + 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( + "AgentCore shell closed by the remote session", + "remote_session_killed", + ) + return line + + 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 self._closed_error( + "AgentCore shell is not running (reader ended); refusing to send " + "to a half-closed transport" + ) + 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 + 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: + """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/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..f26fbcaa3 100644 --- a/src/benchflow/sandbox/providers.py +++ b/src/benchflow/sandbox/providers.py @@ -36,6 +36,14 @@ 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 + #: 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: @@ -46,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", @@ -61,6 +76,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 +107,16 @@ 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 +) + + 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..a77bec38b 100644 --- a/src/benchflow/task/runtime_capabilities.py +++ b/src/benchflow/task/runtime_capabilities.py @@ -15,7 +15,13 @@ 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._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 ( NetworkMode, TaskConfig, @@ -261,11 +267,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: @@ -510,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_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_builder.py b/tests/test_agentcore_builder.py new file mode 100644 index 000000000..5f0c2af82 --- /dev/null +++ b/tests/test_agentcore_builder.py @@ -0,0 +1,799 @@ +"""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 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 + +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) + + +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): + 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_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"]) + + 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"]) + + 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)) + + +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 + + def test_character_classes_are_honoured(self, tmp_path): + """`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") + 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 + + @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. + + Guards PR #937. + """ + 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. + + Guards PR #937. + """ + 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 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 + + 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, *first + ) != provisioning.build_context_digest(tmp_path, *second) + + +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._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 + 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() + 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. + + Guards PR #937. + """ + 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") + + @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 + + @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 new file mode 100644 index 000000000..51ae4571c --- /dev/null +++ b/tests/test_agentcore_parallel.py @@ -0,0 +1,906 @@ +"""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 +import time +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.sandbox.agentcore_image import AgentCoreImagePublisher +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._images._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)._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)._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.""" + 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)._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. + + 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(AgentCoreImagePublisher, "_ensure_ecr_repository"), + patch.object(AgentCoreImagePublisher, "_image_exists", return_value=False), + patch.object(AgentCoreImagePublisher, "_build_and_push", new=_build), + patch.object( + 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._images.publish(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(AgentCoreImagePublisher, "_ensure_ecr_repository"), + patch.object(AgentCoreImagePublisher, "_image_exists", return_value=False), + patch.object(AgentCoreImagePublisher, "_build_and_push", new=_flaky), + patch.object( + 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._images.publish(force_build=False) + await env._images.publish(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 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"} + }, + "lifecycleConfiguration": env._lifecycle_configuration(), + "roleArn": "arn:aws:iam::1:role/rt", + "networkConfiguration": {"networkMode": "PUBLIC"}, + "protocolConfiguration": {"serverProtocol": "HTTP"}, + } + + 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"} + }, + "lifecycleConfiguration": { + "idleRuntimeSessionTimeout": 900, + "maxLifetime": 28800, + }, + "roleArn": "arn:aws:iam::1:role/rt", + "networkConfiguration": {"networkMode": "PUBLIC"}, + "protocolConfiguration": {"serverProtocol": "HTTP"}, + } + + with pytest.raises(SandboxStartupError, match="wrong image"): + AgentCoreSandbox._verify_adopted_runtime( + control, + "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): + """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="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 + 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, + }, + "roleArn": "arn:aws:iam::1:role/rt", + "networkConfiguration": {"networkMode": "PUBLIC"}, + "protocolConfiguration": {"serverProtocol": "HTTP"}, + } + + 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 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.""" + 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) + + +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): + """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() + 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 + 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: + """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. + + Guards PR #937. + """ + 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. + + Guards PR #937. + """ + 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. + + 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. + + 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" + 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"] + + @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): + 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. + + Guards PR #937. + """ + 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. + + 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): + """Guards PR #937 through the deprecated cleanup alias.""" + 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): + """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"): + reap_stale_sandboxes(MagicMock(), dry_run=True, **kwargs) + + @pytest.mark.parametrize("command", ["sandbox", "environment"]) + def test_cli_rejects_negative_max_age(self, command): + """Guards PR #937: both current and deprecated aliases 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 + + +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): + """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 + 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._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), + ): + 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): + """Guards PR #937: generated names are an explicit backend contract. + + 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" + (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 + + @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 new file mode 100644 index 000000000..19c425a3d --- /dev/null +++ b/tests/test_agentcore_sandbox.py @@ -0,0 +1,500 @@ +"""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 contextlib +import os +import re +import shlex +import tarfile +from io import BytesIO +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) + 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: + """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.""" + 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") + + 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): + 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) + + @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_image import PING_SHIM + + return builders.BuildRequest( + context_dir=sandbox.environment_dir, + dockerfile_text=sandbox._images.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.""" + text = sandbox._images.generated_dockerfile_text() + + assert "FROM python:3.12-slim" in text + assert "benchflow_agentcore_shim.py" in text + assert "EXPOSE 8080" in text + + def test_task_dockerfile_is_not_modified_in_place(self, sandbox, tmp_path): + from benchflow.sandbox import agentcore_builder as builders + + 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): + """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)) 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 + + 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()} + staged = None + with contextlib.suppress(RuntimeError): # noqa: SIM117 + 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( + 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"), + ) + + assert "FROM python:3.12-slim" in env._images.generated_dockerfile_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, + 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" + ) + 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_agentcore_transport.py b/tests/test_agentcore_transport.py new file mode 100644 index 000000000..67504fd2b --- /dev/null +++ b/tests/test_agentcore_transport.py @@ -0,0 +1,484 @@ +"""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 +import os +import re +import shlex +import sys +from types import ModuleType, 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"): + 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 + + +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 _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 + 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_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 + 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" + + @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 _runtime_sdk(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_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.""" + 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 ( + _runtime_sdk(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): + """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): + """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( + [ + _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): + """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. + """ + shell = _FakeShell([]) + proc = _process(shell) + await proc._drain_frames() + + 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(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): + """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 ( + _runtime_sdk(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" + + +@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) 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_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/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..259f068a3 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" @@ -334,6 +353,7 @@ deepagents = [ { name = "langchain-openai" }, ] dev = [ + { name = "pathspec" }, { name = "pre-commit" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -345,6 +365,11 @@ judge = [ { name = "google-genai" }, { name = "openai" }, ] +sandbox-agentcore = [ + { name = "bedrock-agentcore" }, + { name = "boto3" }, + { name = "pathspec" }, +] sandbox-daytona = [ { name = "daytona" }, { name = "tenacity" }, @@ -366,7 +391,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" }, @@ -377,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" }, @@ -392,34 +421,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]] @@ -2773,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" @@ -3667,14 +3705,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]]