From f84cdfe20fada7b8df68a61f9cb9116b80560dd6 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 7 Oct 2026 08:24:08 -0700 Subject: [PATCH] refactor(sandbox): object-store names for the sandbox object copies; read handed-in urls from the store that holds them Sandbox.write_file_from_object and VmSandbox.load_object_file replace write_file_from_s3 and load_s3_file, which stay as deprecated aliases. The implementations were already store-neutral; only the names, the s3_url parameters and the /tmp/_s3_ staging prefix said S3. Every core call site moves to the new names. The sites that read, sign or key a url they were handed now resolve the store that holds it (get_object_store_at) instead of the configured one: the VM download and image signing, the Modal VM download, ModalSandbox's object write, the rubrics verifier's trajectory reads and snapshot_agent_state's per-service capture. Under namespace routing outside an @local run, a url the local store holds was signed or keyed by the configured store, which for S3 raised "is not an object in s3://". Sites that mint new urls keep the configured store. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/a2a_agent/a2a_agent.py | 4 +- src/agent_env/env/env.py | 8 ++-- src/agent_env/env/envs/mcp_server.py | 4 +- src/agent_env/env/envs/multi_env.py | 4 +- src/agent_env/env/envs/website.py | 2 +- .../sandbox_providers/modal_sandbox.py | 4 +- .../sandbox_providers/modal_vm_sandbox.py | 4 +- .../providers/sandbox_providers/sandbox.py | 37 +++++++++++----- .../task_step/task_steps/install_agent.py | 2 +- .../task_step/task_steps/load_artifact.py | 14 +++--- .../task_steps/run_docker_container.py | 4 +- .../task_steps/snapshot_agent_state.py | 19 ++++---- .../task_steps/verifiers/rubrics_verifier.py | 13 +++--- .../test_multiturn_verifier_delivery.py | 4 +- .../envs/test_mcp_server_container_mode.py | 14 +++--- .../sandbox_providers/modal_sandbox_test.py | 20 ++++++++- .../modal_vm_sandbox_test.py | 31 ++++++++----- .../sandbox_provider_test.py | 4 +- .../sandbox_providers/vm_sandbox_test.py | 44 ++++++++++++++++--- tst/unit/store/fakes.py | 21 +++++++++ .../task_step/run_docker_container_test.py | 2 +- .../test_load_artifact_container_sandbox.py | 8 ++-- tst/unit/task_step/test_load_artifact_file.py | 8 ++-- .../task_step/test_load_artifact_vm_target.py | 14 +++--- .../task_step/test_snapshot_agent_state.py | 31 +++++++++++-- .../test_rubrics_verifier_multiturn.py | 25 ++++++++--- tst/unit/utils/test_paths.py | 2 +- 27 files changed, 240 insertions(+), 107 deletions(-) diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index e389cf90..a0dc1707 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -241,11 +241,11 @@ async def install_cli(deployed: "DeployedA2AAgent", cli_artifact: "CliArtifact", entrypoint_path = f"{target_root}/{cli_artifact.entrypoint}" env_dir = posixpath.dirname(entrypoint_path) - # write .env first — write_file_from_text creates the parent dir; write_file_from_s3 does not. + # write .env first — write_file_from_text creates the parent dir; write_file_from_object does not. await sandbox.write_file_from_text(f"AGENT_ENV_GATEWAY_URL={gateway_url}\n", f"{env_dir}/.env") for rel_path, file_artifact in cli_artifact.get_cli_files().get_file_artifacts().items(): - await sandbox.write_file_from_s3(file_artifact.object_url, f"{target_root}/{rel_path}") + await sandbox.write_file_from_object(file_artifact.object_url, f"{target_root}/{rel_path}") if isinstance(sandbox, VmSandbox): await sandbox.exec_script(f"docker exec -u 0 {shlex.quote(sandbox.container_name)} chmod +x {entrypoint_path}") diff --git a/src/agent_env/env/env.py b/src/agent_env/env/env.py index f083b9a8..02b97788 100644 --- a/src/agent_env/env/env.py +++ b/src/agent_env/env/env.py @@ -326,9 +326,9 @@ async def load_file_artifact_universe( Files land at ``/``. Nested relative paths (e.g. ``subdir/file.txt``) are supported — parent dirs are created - as needed. Each file is fetched directly from S3 by the sandbox via - a presigned URL, so we never round-trip the bytes through this - process. + as needed. The sandbox fetches each file directly from the object + store through a signed URL where the store can sign one, so the bytes + don't round-trip through this process. Default implementation uses ``self._sandbox`` (set by ``deploy()`` / ``from_deployed_env()`` on most env subclasses). Subclasses that @@ -373,7 +373,7 @@ async def load_file_artifact_universe( for idx, (filename, file_artifact) in enumerate(file_artifacts.items(), 1): dest_path = loaded[filename] logger.info(f" [{idx}/{total}] {file_artifact.object_url} -> {dest_path}") - await sandbox.load_s3_file(file_artifact.object_url, dest_path) + await sandbox.load_object_file(file_artifact.object_url, dest_path) logger.info( f"Loaded FileArtifactUniverse '{file_artifact_universe.id}' " diff --git a/src/agent_env/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index 5dff7545..238db5aa 100644 --- a/src/agent_env/env/envs/mcp_server.py +++ b/src/agent_env/env/envs/mcp_server.py @@ -226,7 +226,7 @@ async def _staged_artifact_size(self, file_artifact) -> int | None: async def _copy_artifact_into_container(self, file_artifact) -> str: container_path = f"/data/{file_artifact.filename}" if self._sandbox.mode == SANDBOX_MODE_CONTAINER: - await self._sandbox.write_file_from_s3(file_artifact.object_url, container_path) + await self._sandbox.write_file_from_object(file_artifact.object_url, container_path) else: # Stage on the disk-backed app dir, not /tmp (tmpfs/RAM): a large # artifact (github ~8GB) OOM-kills the copy under concurrent loads. @@ -235,7 +235,7 @@ async def _copy_artifact_into_container(self, file_artifact) -> str: stage_dir = f"{GATEWAY_APP_DIR}/_artifact_staging" vm_temp_path = f"{stage_dir}/{self.environment_name}-{uuid.uuid4().hex[:8]}-{file_artifact.filename}" await self._sandbox.exec_script(f"mkdir -p {stage_dir}") - await self._sandbox.load_s3_file(file_artifact.object_url, vm_temp_path) + await self._sandbox.load_object_file(file_artifact.object_url, vm_temp_path) container_id = await self._env_provider._get_container_id(self._sandbox, self.environment_name) await self._sandbox.exec_script(f"docker exec {container_id} mkdir -p /data") await self._sandbox.docker_cp(vm_temp_path, f"{container_id}:{container_path}") diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 8c4cf81b..abba5865 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -353,9 +353,9 @@ async def _load_one(idx: int, artifact): dest_path = f"{dest_dir}/{file_artifact.filename}" logger.info(f"Downloading metadata '{key}' -> {dest_path}") if self._sandbox.mode == SANDBOX_MODE_CONTAINER: - await self._sandbox.write_file_from_s3(file_artifact.object_url, dest_path) + await self._sandbox.write_file_from_object(file_artifact.object_url, dest_path) else: - await self._sandbox.load_s3_file(file_artifact.object_url, dest_path) + await self._sandbox.load_object_file(file_artifact.object_url, dest_path) metadata_filepaths[key] = dest_path logger.info(f"Successfully downloaded metadata '{key}'") return LoadEnvironmentUniverseArtifactResult( diff --git a/src/agent_env/env/envs/website.py b/src/agent_env/env/envs/website.py index 7f316f04..5d693bb4 100644 --- a/src/agent_env/env/envs/website.py +++ b/src/agent_env/env/envs/website.py @@ -185,7 +185,7 @@ async def _copy_artifact_into_container(self, file_artifact) -> str: container_path = f"/tmp/data/{filename}" compose = f"docker compose -f {DOCKER_COMPOSE_PATH}" - await self._sandbox.load_s3_file(file_artifact.object_url, vm_path) + await self._sandbox.load_object_file(file_artifact.object_url, vm_path) await self._sandbox.exec_script(f"{compose} exec -T {backend_service} mkdir -p /tmp/data") await self._sandbox.exec_script(f"{compose} cp {vm_path} {backend_service}:{container_path}") await self._sandbox.exec_script(f"rm -f {vm_path}") diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index 6b6abf07..ea14c39c 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -152,8 +152,8 @@ async def _write_stream_via_exec( stderr = await proc.stderr.read.aio() raise RuntimeError(f"write to {destination_path} failed (exit={rc}): {stderr!r}") - async def write_file_from_s3(self, s3_url: str, destination_path: str) -> None: - body = await asyncio.to_thread(get_config().get_object_store().open, s3_url) + async def write_file_from_object(self, object_url: str, destination_path: str) -> None: + body = await asyncio.to_thread(get_config().get_object_store_at(object_url).open, object_url) with contextlib.closing(body): async def _aiter() -> AsyncIterator[bytes]: diff --git a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py index 107281c7..7e2ca1b5 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -105,7 +105,7 @@ class ModalVmSandbox(VmSandbox): """A Modal VM sandbox (vm_runtime) presented through the VmSandbox contract. - Subclasses VmSandbox so it inherits exec_script / load_docker_images / load_s3_file / + Subclasses VmSandbox so it inherits exec_script / load_docker_images / load_object_file / write_file_* unchanged — they are built on exec_with_output -> exec + an in-VM Docker daemon, which this class provides. Only the Modal-specific pieces (raw exec, terminate, dockerd startup) are implemented here. @@ -194,7 +194,7 @@ async def _ensure_dockerd_running(self) -> None: async def _download_object_to_vm(self, object_url: str, vm_path: str) -> None: """Presigned objects go through aria2c (see _DL_*); a store that cannot presign keeps the base path.""" - store = get_config().get_object_store() + store = get_config().get_object_store_at(object_url) signed = await asyncio.to_thread(store.signed_get_url, object_url) signed_at = time.monotonic() if signed is None: diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 4d0257aa..4da644b8 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -11,6 +11,7 @@ import tempfile import time import uuid +import warnings import weakref from abc import ABC, abstractmethod from dataclasses import dataclass, replace @@ -162,9 +163,17 @@ async def exec_with_output(self, *args: str) -> tuple[int, str, str]: exit_code = await process.wait() return exit_code, stdout.decode(), stderr.decode() + async def write_file_from_object(self, object_url: str, destination_path: str) -> None: + """Write the object store's object at ``object_url`` into the agent process's filesystem at + ``destination_path``.""" + raise NotImplementedError(f"{self.__class__.__name__} does not support write_file_from_object") + async def write_file_from_s3(self, s3_url: str, destination_path: str) -> None: - """Write a file from S3 into the agent process's filesystem at destination_path.""" - raise NotImplementedError(f"{self.__class__.__name__} does not support write_file_from_s3") + """Deprecated: ``write_file_from_object``.""" + warnings.warn( + "Sandbox.write_file_from_s3 is deprecated; use write_file_from_object", DeprecationWarning, stacklevel=2 + ) + await self.write_file_from_object(s3_url, destination_path) async def write_file_from_url(self, url: str, destination_path: str) -> None: """Download an HTTP(S) URL into the agent process's filesystem at destination_path.""" @@ -269,7 +278,7 @@ async def _signed_image_urls(artifacts: list) -> list[str | None]: """Each image tarball's signed URL, or None where the store cannot sign one. Signed concurrently and off the event loop, since a remote signer is a network round trip.""" logger.info(f"Loading {len(artifacts)} Docker image(s) into the sandbox...") - object_store = get_config().get_object_store() + config = get_config() slots = _sign_slots.setdefault(asyncio.get_running_loop(), asyncio.Semaphore(_CONCURRENT_SIGNS)) failed = False @@ -279,7 +288,8 @@ async def sign(artifact) -> str | None: if failed: return None try: - return await asyncio.to_thread(object_store.signed_get_url, artifact.tar_gz_object_url) + store = config.get_object_store_at(artifact.tar_gz_object_url) + return await asyncio.to_thread(store.signed_get_url, artifact.tar_gz_object_url) except BaseException: failed = True raise @@ -315,13 +325,18 @@ async def _load_docker_images(self, artifacts: list, signed_urls: list[str | Non raise RuntimeError(f"{artifact.image_name} image not found. stdout: {stdout}") logger.info(" All images loaded successfully") + async def load_object_file(self, object_url: str, destination_path: str) -> None: + """Download the object store's object at ``object_url`` onto the VM host at ``destination_path``.""" + await self._download_object_to_vm(object_url, destination_path) + async def load_s3_file(self, s3_url: str, destination_path: str) -> None: - """Download an object from the object store into the sandbox.""" - await self._download_object_to_vm(s3_url, destination_path) + """Deprecated: ``load_object_file``.""" + warnings.warn("VmSandbox.load_s3_file is deprecated; use load_object_file", DeprecationWarning, stacklevel=2) + await self.load_object_file(s3_url, destination_path) async def _download_object_to_vm(self, object_url: str, vm_path: str) -> None: """Place object_url onto the VM host at vm_path, backend-agnostically.""" - object_store = get_config().get_object_store() + object_store = get_config().get_object_store_at(object_url) signed = await asyncio.to_thread(object_store.signed_get_url, object_url) if signed is not None: await self.exec_script(f"curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(signed)} -o {shlex.quote(vm_path)}") @@ -361,10 +376,10 @@ def _staging_path(kind: str, destination_path: str) -> str: host's /tmp, as local ones do.""" return f"/tmp/_{kind}_{uuid.uuid4().hex[:12]}_{destination_path.replace('/', '_').lstrip('_')}" - async def write_file_from_s3(self, s3_url: str, destination_path: str) -> None: - vm_path = self._staging_path("s3", destination_path) + async def write_file_from_object(self, object_url: str, destination_path: str) -> None: + vm_path = self._staging_path("obj", destination_path) try: - await self.load_s3_file(s3_url, vm_path) + await self.load_object_file(object_url, vm_path) await self._copy_into_container(vm_path, destination_path) finally: await self._remove_vm_temp_file(vm_path) @@ -438,7 +453,7 @@ async def stage_files_into_container(sandbox: Sandbox, file_artifacts: dict[str, async def _stage(filename: str, file_artifact: Any) -> None: async with sem: logger.info(f" {file_artifact.object_url} -> {loaded[filename]}") - await sandbox.write_file_from_s3(file_artifact.object_url, loaded[filename]) + await sandbox.write_file_from_object(file_artifact.object_url, loaded[filename]) await asyncio.gather(*(_stage(fn, fa) for fn, fa in file_artifacts.items())) return loaded diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index 4677aa0c..659f6e31 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -227,7 +227,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: agent_ctx_tar = f"{work_dir}/agent-ctx.tar.gz" await sandbox.exec_script(f"rm -rf {shlex.quote(work_dir)} && mkdir -p {shlex.quote(work_dir)}") logger.info(f"Downloading agent build context {agent.docker_image_artifact.build_context_object_url} -> {agent_ctx_tar}") - await sandbox.load_s3_file(agent.docker_image_artifact.build_context_object_url, agent_ctx_tar) + await sandbox.load_object_file(agent.docker_image_artifact.build_context_object_url, agent_ctx_tar) # Resolve each required_param via the registry, shlex.quote so str.format # substitution produces shell-safe tokens. diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index 6bc26a1c..b1564c83 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -264,7 +264,7 @@ def _resolve_inputs(self, context: TaskStepContext) -> _LoadInputs: return _LoadInputs(artifacts=resolved_artifacts, urls=urls, destination_path=destination_path) async def _load_url_onto_vm(self, sandbox, url: str, destination_path: str) -> None: - """Host counterpart of ``sandbox.load_s3_file``; ``write_file_from_url`` targets a container instead.""" + """Host counterpart of ``sandbox.load_object_file``; ``write_file_from_url`` targets a container instead.""" from agent_env.providers.sandbox_providers.sandbox import CURL_RETRY_FLAGS parent = posixpath.dirname(destination_path) @@ -643,7 +643,7 @@ async def _stage_environment_payload_into_container( if isinstance(sandbox, VmSandbox): vm_payload = f"/tmp/_svc_{token}_{safe_name}" vm_stage = f"/tmp/_svc_stage_{token}" - await sandbox.load_s3_file(file_artifact.object_url, vm_payload) + await sandbox.load_object_file(file_artifact.object_url, vm_payload) try: # exec_script raises on a non-zero exit, but only with the raw # script output — re-raise with the artifact context so a failed @@ -679,7 +679,7 @@ async def _stage_environment_payload_into_container( # Container-mode sandbox: the sandbox IS the agent container — expand # in place at the destination. payload_tmp = f"/tmp/_svc_{token}_{safe_name}" - await sandbox.write_file_from_s3(file_artifact.object_url, payload_tmp) + await sandbox.write_file_from_object(file_artifact.object_url, payload_tmp) try: exit_code, out, err = await sandbox.exec_with_output( "python3", "-c", _ENVIRONMENT_PAYLOAD_EXPAND_SCRIPT, payload_tmp, destination @@ -712,7 +712,7 @@ async def _stage_environment_payload_into_container( async def _load_universe_onto_vm(sandbox, universe, destination: str) -> list[str]: - """Pull each file in `universe` from S3 straight onto the VM host — no temp file, no copy inward.""" + """Pull each file in `universe` from the object store straight onto the VM host — no temp file, no copy inward.""" file_artifacts = universe.get_file_artifacts() if not file_artifacts: logger.warning( @@ -734,7 +734,7 @@ async def _load_universe_onto_vm(sandbox, universe, destination: str) -> list[st logger.info(f" [{idx}/{total}] {fa.object_url} -> {dest_path}") if parent and parent != destination: await sandbox.exec_script(f"mkdir -p {shlex.quote(parent)}") - await sandbox.load_s3_file(fa.object_url, dest_path) + await sandbox.load_object_file(fa.object_url, dest_path) loaded.append(filename) return loaded @@ -757,7 +757,7 @@ async def _load_universe_into_sandbox_container(sandbox, universe, destination: async def _load_universe_into_container(sandbox, container_name: str, universe, destination: str) -> list[str]: - """For each file in `universe`: pull from S3 to a VM temp path, then `docker cp` into `container_name` at `destination/`.""" + """For each file in `universe`: pull from the object store to a VM temp path, then `docker cp` into `container_name` at `destination/`.""" file_artifacts = universe.get_file_artifacts() if not file_artifacts: logger.warning( @@ -780,7 +780,7 @@ async def _load_universe_into_container(sandbox, container_name: str, universe, parent = posixpath.dirname(dest_path) vm_temp = f"/tmp/_load_{uuid.uuid4().hex[:8]}_{filename.replace('/', '_')}" logger.info(f" [{idx}/{total}] {fa.object_url} -> {container_name}:{dest_path}") - await sandbox.load_s3_file(fa.object_url, vm_temp) + await sandbox.load_object_file(fa.object_url, vm_temp) if parent and parent != destination: await sandbox.exec_script( f"docker exec {shlex.quote(container_name)} mkdir -p {shlex.quote(parent)}" diff --git a/src/agent_env/task_step/task_steps/run_docker_container.py b/src/agent_env/task_step/task_steps/run_docker_container.py index b73c89e7..9c8013d0 100644 --- a/src/agent_env/task_step/task_steps/run_docker_container.py +++ b/src/agent_env/task_step/task_steps/run_docker_container.py @@ -341,7 +341,7 @@ async def _stage_from_universe(self, sandbox, sandbox_id: str, work_dir: str) -> for idx, (filename, fa) in enumerate(file_artifacts.items(), 1): dest_path = loaded[filename] logger.info(f" [{idx}/{total}] {fa.object_url} -> {dest_path}") - await sandbox.load_s3_file(fa.object_url, dest_path) + await sandbox.load_object_file(fa.object_url, dest_path) @staticmethod async def _stage_zip_from_url(sandbox, url: str, work_dir: str) -> None: @@ -369,7 +369,7 @@ async def _stage_zip_from_url(sandbox, url: str, work_dir: str) -> None: archive_path = posixpath.join(work_dir, "_context.zip") logger.info(f"Downloading docker context from {url} into {work_dir}") if from_store: - await sandbox.load_s3_file(url, archive_path) + await sandbox.load_object_file(url, archive_path) else: await sandbox.exec_script( f"curl -fsSL {shlex.quote(url)} -o {shlex.quote(archive_path)}" diff --git a/src/agent_env/task_step/task_steps/snapshot_agent_state.py b/src/agent_env/task_step/task_steps/snapshot_agent_state.py index 25b1b7a3..0dccc449 100644 --- a/src/agent_env/task_step/task_steps/snapshot_agent_state.py +++ b/src/agent_env/task_step/task_steps/snapshot_agent_state.py @@ -21,15 +21,15 @@ class SnapshotAgentStateTaskStep(TaskStep): Looks up the deployed agent (by ``agent_name``) and the source ``context_id`` via a previously-recorded ``PromptResponse`` (by ``prompt_id``). Calls the agent's ``/ext/snapshot`` extension to write the conversation transcript to - S3, then wraps that S3 prefix as a ``FileArtifactUniverse`` via + the object store, then wraps that prefix as a ``FileArtifactUniverse`` via ``put_existing`` (no download/re-upload). Appends an entry recording ``{id, version, bundle_object_url, source_agent_name, source_context_id}`` to ``context.metadata['agent_snapshots']`` (a list of such entries). - When ``env_id`` is set, the step ALSO dumps each MCP service's current - state by hitting ``{gateway_url}/svc/mcp-{service}/export-state`` on the - deployed env, uploads each as JSON alongside the workspace tarball, and - writes a JSON-stringified ``{service: s3_url}`` (unsigned ``s3://`` object + When ``env_id`` is set, the step ALSO reads each MCP service's current + state from the deployed env (its ``data/get`` JSON, else its + ``/export-state``), uploads each as JSON alongside the workspace tarball, and + writes a JSON-stringified ``{service: object_url}`` (unsigned object-store references) to ``context.metadata['snapshot_json_url']``. A downstream verifier re-signs each value before handing the map to its test script, so the script can fetch each service's state and assert against it. Without this, only the @@ -133,7 +133,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # universe_artifact_id are both configured. Mirrors the export-state # pattern in `multienv_validator/verify_universe_roundtrip.py:_export_all`. # Result lands in `context.metadata['snapshot_json_url']` as a - # JSON-stringified `{service: s3_url}` (unsigned `s3://` references). + # JSON-stringified `{service: object_url}` (unsigned object-store references). # A downstream verifier re-signs each value and hands the map to its test # script, which accepts either a single URL or this map. if self.env_id and self.universe_artifact_id: @@ -144,7 +144,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: async def _capture_universe_state( self, context: TaskStepContext, capture_prefix: str ) -> None: - """Read each service's state as JSON, upload it, publish `s3://` refs.""" + """Read each service's state as JSON, upload it beside the workspace capture, and publish the + object URLs.""" from agent_env.artifact import EnvironmentUniverseArtifact from agent_env.config import get_config @@ -170,9 +171,9 @@ async def _capture_universe_state( ) return - # Reuse the workspace capture's S3 prefix — services land in a + # Reuse the workspace capture's prefix, in the store that holds it — services land in a # `services/` subdir alongside the workspace tarball. - store = get_config().get_object_store() + store = get_config().get_object_store_at(capture_prefix) key_prefix = store.get_object_key(capture_prefix).rstrip("/") + "/" from agent_env.env import legacy_protocol diff --git a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py index 48e8a2a1..e0b44bc2 100644 --- a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py @@ -586,7 +586,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if multiturn: await self._load_per_turn_trajectories(sandbox, per_turn_uris, multiturn_dir, filter_to_apply) else: - await sandbox.write_file_from_s3(compacted.container_s3_uri, self.TRAJECTORY_CONTAINER_PATH) + await sandbox.write_file_from_object(compacted.container_s3_uri, self.TRAJECTORY_CONTAINER_PATH) if compacted.tool_result_files: await self._load_tool_result_files(sandbox, compacted.tool_result_files) compacted.tool_result_files = [] @@ -816,11 +816,10 @@ def _filter_trajectory(self, s3_uri: str, trajectory_filter: TrajectoryFilter) - from agent_env.config import get_config config = get_config() - object_store = config.get_object_store() - spans = json.loads(object_store.get(s3_uri)) + spans = json.loads(config.get_object_store_at(s3_uri).get(s3_uri)) filtered, tool_result_files = compact_otel_trajectory(spans, trajectory_filter) compact_key = f"{config.get_artifact_key_prefix()}compacted-trajectories/{uuid.uuid4().hex}.json" - object_url = object_store.put(compact_key, json.dumps(filtered).encode(), content_type="application/json") + object_url = config.get_object_store().put(compact_key, json.dumps(filtered).encode(), content_type="application/json") return object_url, tool_result_files async def _load_tool_result_files(self, sandbox, tool_result_files: list[tuple[str, str]]) -> None: @@ -833,11 +832,11 @@ async def _load_trajectory_into_container(self, sandbox_id: str, s3_uri: str) -> """Load trajectory from S3 into an existing agent's container.""" from agent_env.providers.sandbox_providers.sandbox_provider import get_agent_sandbox_provider sandbox = await get_agent_sandbox_provider().get_sandbox(sandbox_id) - await sandbox.write_file_from_s3(s3_uri, self.TRAJECTORY_CONTAINER_PATH) + await sandbox.write_file_from_object(s3_uri, self.TRAJECTORY_CONTAINER_PATH) def _read_trajectory_text(self, s3_uri: str) -> str: from agent_env.config import get_config - return get_config().get_object_store().get(s3_uri).decode("utf-8", errors="replace") + return get_config().get_object_store_at(s3_uri).get(s3_uri).decode("utf-8", errors="replace") async def _load_per_turn_trajectories( self, sandbox, per_turn_uris: list[str], container_dir: str, @@ -860,7 +859,7 @@ async def _load_per_turn_trajectories( await sandbox.write_file_from_text(content, dest) tool_result_files.extend(turn_files) else: - await sandbox.write_file_from_s3(uri, dest) + await sandbox.write_file_from_object(uri, dest) if tool_result_files: await self._load_tool_result_files(sandbox, tool_result_files) logger.info( diff --git a/tst/integration/task_step/test_multiturn_verifier_delivery.py b/tst/integration/task_step/test_multiturn_verifier_delivery.py index 85203500..b6593b8d 100644 --- a/tst/integration/task_step/test_multiturn_verifier_delivery.py +++ b/tst/integration/task_step/test_multiturn_verifier_delivery.py @@ -6,7 +6,7 @@ they are intact and parseable *from inside* the sandbox (the deterministic ceiling for "loaded by the agent"). -Uses a **bare Modal sandbox**: ``ModalSandbox.write_file_from_s3`` streams into +Uses a **bare Modal sandbox**: ``ModalSandbox.write_file_from_object`` streams into the sandbox FS directly, so no agent needs to be deployed. The VM-provider path (``docker cp … agent-api:…`` + the container ``mkdir -p`` fix) requires a deployed agent's container and is covered by the full-agent e2e / manual runs. @@ -62,7 +62,7 @@ def _tool_span(tool: str, output) -> dict: def per_turn_uris(): """Upload two per-turn OTel-span trajectories through the configured object store. - The verifier hands them to the sandbox as S3 objects (``Sandbox.write_file_from_s3``), so a store + The verifier hands them to the sandbox as S3 objects (``Sandbox.write_file_from_object``), so a store that mints another scheme fails here rather than inside the container. The objects stay behind: the ObjectStore API has no delete, and the prefix is unique per run.""" store = get_config().get_object_store() diff --git a/tst/unit/env/envs/test_mcp_server_container_mode.py b/tst/unit/env/envs/test_mcp_server_container_mode.py index b7449d37..635fa0f6 100644 --- a/tst/unit/env/envs/test_mcp_server_container_mode.py +++ b/tst/unit/env/envs/test_mcp_server_container_mode.py @@ -62,8 +62,8 @@ async def test_a_container_deploy_stages_into_the_server_and_records_every_conta assert (deployed.sandbox_id, deployed.sandbox_type) == ("gw", "modal") assert deployed.sandbox_ids == _RECORDED assert {i for i, _ in _env_sandbox_ids(deployed)} == {sb.sandbox_id for sb in env._env_provider._container_sandboxes} - sandboxes["srv"].write_file_from_s3.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") - sandboxes["gw"].write_file_from_s3.assert_not_awaited() + sandboxes["srv"].write_file_from_object.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") + sandboxes["gw"].write_file_from_object.assert_not_awaited() @pytest.mark.asyncio @@ -128,8 +128,8 @@ async def test_a_restored_container_env_stages_into_the_server_and_closes_every_ await env.close() assert env._sandbox is None - reattached["srv"].write_file_from_s3.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") - reattached["gw"].write_file_from_s3.assert_not_awaited() + reattached["srv"].write_file_from_object.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") + reattached["gw"].write_file_from_object.assert_not_awaited() assert [name for name, sb in reattached.items() if sb.terminate.await_count == 0] == [] @@ -260,7 +260,7 @@ async def test_a_deploy_without_a_gateway_loads_straight_into_the_server(sent): await _deploy(env, sandboxes) await env.load_environment_artifact(_artifact()) - sandboxes["srv"].write_file_from_s3.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") + sandboxes["srv"].write_file_from_object.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") assert {str(r.url) for r in sent} == {f"{_SRV}/agentenv"} @@ -275,7 +275,7 @@ async def test_a_restored_env_without_a_gateway_loads_into_the_server_and_closes await env.load_environment_artifact(_artifact()) await env.close() - reattached["srv"].write_file_from_s3.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") + reattached["srv"].write_file_from_object.assert_awaited_once_with("s3://bucket/email.json", "/data/email.json") assert {str(r.url) for r in sent} == {f"{_SRV}/agentenv"} and reattached["srv"].mode == SANDBOX_MODE_CONTAINER assert reattached["srv"].terminate.await_count > 0 @@ -312,7 +312,7 @@ def _env(env_provider_type: str = "gateway") -> MCPServerEnv: def _sandbox(sandbox_id: str, mode: str = SANDBOX_MODE_CONTAINER, sandbox_type: str = "modal") -> MagicMock: - return MagicMock(sandbox_id=sandbox_id, mode=mode, type=sandbox_type, terminate=AsyncMock(), write_file_from_s3=AsyncMock()) + return MagicMock(sandbox_id=sandbox_id, mode=mode, type=sandbox_type, terminate=AsyncMock(), write_file_from_object=AsyncMock()) def _containers() -> dict[str, MagicMock]: diff --git a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py index a095df55..2a3ffb4e 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py @@ -6,7 +6,10 @@ import pytest -from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider +from agent_env.config import get_config +from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandbox, ModalSandboxProvider +from agent_env.store import set_object_store +from tst.unit.store.fakes import ConfiguredObjectStore class _BareSandboxTimeoutError(Exception): @@ -245,3 +248,18 @@ async def test_vnc_port_controls_the_tunnel_set_and_the_vnc_url( assert mock_create.aio.call_args.kwargs["encrypted_ports"] == expected_ports assert sandbox.vnc_url == expected_vnc_url + + +@pytest.mark.asyncio +async def test_write_file_from_object_reads_a_url_from_the_store_that_holds_it(cli_routing): + set_object_store(ConfiguredObjectStore()) + url = get_config().get_object_store_for("@local/~/t").put("objects/a.json", b"LOCAL") + sandbox = ModalSandbox(MagicMock(), tunnel_urls={}) + written = [] + + async def collect(chunks, destination_path): + written.append((b"".join([chunk async for chunk in chunks]), destination_path)) + + sandbox._write_stream_via_exec = collect + await sandbox.write_file_from_object(url, "/data/a.json") + assert written == [(b"LOCAL", "/data/a.json")] diff --git a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py index c315f7c6..afa83046 100644 --- a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py @@ -7,6 +7,7 @@ from __future__ import annotations +import base64 from unittest.mock import AsyncMock, MagicMock import pytest @@ -130,9 +131,10 @@ def test_provider_init_is_self_contained(): # ── object downloads: the aria2c range pool ─────────────────────────────────────── -from agent_env.config import reset_config # noqa: E402 +from agent_env.config import get_config, reset_config # noqa: E402 from agent_env.providers.sandbox_providers import modal_vm_sandbox as mvs # noqa: E402 from agent_env.store import set_object_store # noqa: E402 +from tst.unit.store.fakes import ConfiguredObjectStore # noqa: E402 URL = "https://bucket.s3.amazonaws.com/artifacts/file/x/1/x.zip?X-Amz-Signature=abc&X-Amz-Expires=3600" @@ -230,7 +232,7 @@ def test_parse_download_summary(): async def test_download_first_attempt_ok(store): """The URL signed to pick the path is the first attempt's.""" vm = _ScriptedVm([_summary()]) - await vm.load_s3_file("s3://b/x.zip", "/app/x.zip") + await vm.load_object_file("s3://b/x.zip", "/app/x.zip") assert len(vm.downloads()) == 1 assert store.minted == 1 @@ -238,7 +240,7 @@ async def test_download_first_attempt_ok(store): @pytest.mark.asyncio async def test_download_retries_with_a_fresh_url(store): vm = _ScriptedVm([_summary(rc=22, msg="status=403 "), _summary()]) - await vm.load_s3_file("s3://b/x.zip", "/app/x.zip") + await vm.load_object_file("s3://b/x.zip", "/app/x.zip") first, second = vm.downloads() assert "?n=1" in first and "?n=2" in second assert 'rm -f "$out" "$out.aria2"' in first and 'rm -f "$out"' not in second # attempt 2 resumes @@ -249,7 +251,7 @@ async def test_a_stale_probe_url_is_re_signed(store, monkeypatch): """A download that queued for a slot can outlive the URL signed to pick its path.""" monkeypatch.setattr(mvs, "_FRESH_URL_SECONDS", -1) vm = _ScriptedVm([_summary()]) - await vm.load_s3_file("s3://b/x.zip", "/app/x.zip") + await vm.load_object_file("s3://b/x.zip", "/app/x.zip") assert "?n=2" in vm.downloads()[0] @@ -260,7 +262,7 @@ async def test_the_first_script_to_run_clears_stale_state_even_after_a_failed_si signs = iter([store.signed_get_url, _refuse, store.signed_get_url]) monkeypatch.setattr(store, "signed_get_url", lambda url, expires_in=3600: next(signs)(url)) vm = _ScriptedVm([_summary()]) - await vm.load_s3_file("s3://b/x.zip", "/app/x.zip") + await vm.load_object_file("s3://b/x.zip", "/app/x.zip") [first] = vm.downloads() assert 'rm -f "$out" "$out.aria2"' in first @@ -270,7 +272,7 @@ async def test_a_failed_re_sign_costs_one_attempt(store, monkeypatch): signs = iter([store.signed_get_url, _refuse, store.signed_get_url]) monkeypatch.setattr(store, "signed_get_url", lambda url, expires_in=3600: next(signs)(url)) vm = _ScriptedVm([_summary(rc=1, msg="recv failure"), _summary()]) - await vm.load_s3_file("s3://b/x.zip", "/app/x.zip") + await vm.load_object_file("s3://b/x.zip", "/app/x.zip") assert len(vm.downloads()) == 2 @@ -282,14 +284,14 @@ def _refuse(url): async def test_download_not_found_raises_immediately(store): vm = _ScriptedVm([_summary(rc=3, msg="Resource not found"), _summary()]) with pytest.raises(RuntimeError, match="not found"): - await vm.load_s3_file("s3://b/missing.zip", "/app/missing.zip") + await vm.load_object_file("s3://b/missing.zip", "/app/missing.zip") assert len(vm.downloads()) == 1 @pytest.mark.asyncio async def test_download_transport_drop_then_deadline_then_ok(store): vm = _ScriptedVm([-1, -1, _summary(rc=124, nbytes=1000), _summary()]) # exec_script retries -1 once itself - await vm.load_s3_file("s3://b/x.zip", "/app/x.zip") + await vm.load_object_file("s3://b/x.zip", "/app/x.zip") assert len(vm.downloads()) == 4 @@ -298,7 +300,7 @@ async def test_download_attempts_exhausted_cleans_up_and_raises(store, monkeypat monkeypatch.setenv("AGENT_ENV_MODAL_VM_DOWNLOAD_ATTEMPTS", "2") vm = _ScriptedVm([_summary(rc=1, msg="recv failure"), _summary(rc=1, msg="recv failure again")]) with pytest.raises(RuntimeError, match="failed after 2 attempts.*recv failure again"): - await vm.load_s3_file("s3://b/x.zip", "/app/x.zip") + await vm.load_object_file("s3://b/x.zip", "/app/x.zip") assert any(s.startswith("rm -f /app/x.zip") for s in vm.scripts) assert any(s.startswith("rm -f /app/x.zip.aria2") for s in vm.scripts) @@ -315,12 +317,21 @@ def get(self, object_url): set_object_store(_Local()) try: vm = _ScriptedVm([]) - await vm.load_s3_file("file:///store/x", "/app/x") + await vm.load_object_file("file:///store/x", "/app/x") finally: reset_config() assert not vm.downloads() and any("base64 -d" in s for s in vm.scripts) +@pytest.mark.asyncio +async def test_a_url_the_local_store_holds_is_read_from_it_not_the_configured_store(cli_routing): + set_object_store(ConfiguredObjectStore()) + url = get_config().get_object_store_for("@local/~/t").put("objects/x.zip", b"LOCAL") + vm = _ScriptedVm([]) + await vm.load_object_file(url, "/app/x.zip") + assert not vm.downloads() and any(base64.b64encode(b"LOCAL").decode() in s for s in vm.scripts) + + @pytest.mark.asyncio async def test_host_file_writes_stay_under_modals_exec_command_limit(): import base64 diff --git a/tst/unit/providers/sandbox_providers/sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/sandbox_provider_test.py index 2457527a..2992fb1c 100644 --- a/tst/unit/providers/sandbox_providers/sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/sandbox_provider_test.py @@ -320,7 +320,7 @@ async def exec_with_output(self, *args): @pytest.mark.asyncio -async def test_load_s3_file_curl_retries_dns_failures(): +async def test_load_object_file_curl_retries_dns_failures(): """The in-VM S3 download must use --retry-all-errors so a transient `curl: (6) Could not resolve host` is retried locally (curl does not retry exit 6 by default, and --retry-connrefused does not cover it).""" @@ -334,7 +334,7 @@ def signed_get_url(self, url, expires_in=3600): set_object_store(_Signing()) try: vm = _RecordingVm() - await vm.load_s3_file("s3://artifact-bucket/foo/bar.tar", "/tmp/bar.tar") + await vm.load_object_file("s3://artifact-bucket/foo/bar.tar", "/tmp/bar.tar") finally: reset_config() diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py index 75ae6d81..b8f9c313 100644 --- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py @@ -1,4 +1,5 @@ import asyncio +import base64 import threading import time from concurrent.futures import ThreadPoolExecutor @@ -9,7 +10,8 @@ from agent_env.providers.sandbox_providers import sandbox as sandbox_module from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.store import set_object_store -from agent_env.config import reset_config +from agent_env.config import get_config, reset_config +from tst.unit.store.fakes import ConfiguredObjectStore from tst.util.exec_scripts import script_run @@ -190,11 +192,11 @@ async def exec_with_output(self, *args): @pytest.mark.asyncio -async def test_write_file_from_s3_survives_cleanup_blip(signing_store): +async def test_write_file_from_object_survives_cleanup_blip(signing_store): """A transient exec failure on the trailing `rm -f` must not fail the write — the file is already delivered into the container.""" sandbox = _CleanupFailingSandbox() - await sandbox.write_file_from_s3("s3://bucket/data.json", "/work/data.json") + await sandbox.write_file_from_object("s3://bucket/data.json", "/work/data.json") # Delivery happened, and the cleanup was attempted despite the blip. assert any("docker cp" in s for s in sandbox.scripts) assert any(s.startswith("rm -f") for s in sandbox.scripts) @@ -217,11 +219,11 @@ async def test_write_file_from_text_survives_cleanup_blip(): @pytest.mark.asyncio -async def test_write_file_from_s3_still_fails_on_delivery_error(signing_store): +async def test_write_file_from_object_still_fails_on_delivery_error(signing_store): """Failures on the actual delivery (docker cp) must still raise.""" sandbox = _CleanupFailingSandbox(fail_on="docker cp") with pytest.raises(RuntimeError): - await sandbox.write_file_from_s3("s3://bucket/data.json", "/work/data.json") + await sandbox.write_file_from_object("s3://bucket/data.json", "/work/data.json") @pytest.mark.asyncio @@ -282,7 +284,7 @@ async def test_the_store_is_called_off_the_event_loop(signs): sandbox = _RecordingVmSandbox(images_stdout="a\nb\n") artifacts = [SimpleNamespace(tar_gz_object_url=f"s3://bucket/{n}.tar.gz", image_name=f"{n}:1") for n in "ab"] await sandbox.load_docker_images(artifacts) - await sandbox.load_s3_file("s3://bucket/data.json", "/tmp/data.json") + await sandbox.load_object_file("s3://bucket/data.json", "/tmp/data.json") finally: reset_config() assert store.on_loop and not any(store.on_loop) @@ -290,7 +292,7 @@ async def test_the_store_is_called_off_the_event_loop(signs): _WRITES = { "text": lambda sandbox, dest: sandbox.write_file_from_text("body", dest), - "object": lambda sandbox, dest: sandbox.write_file_from_s3("s3://bucket/body.json", dest), + "object": lambda sandbox, dest: sandbox.write_file_from_object("s3://bucket/body.json", dest), "url": lambda sandbox, dest: sandbox.write_file_from_url("https://example.com/body.json", dest), } @@ -422,3 +424,31 @@ async def test_docker_cp_passes_its_paths_as_arguments(remove_source, script): async def test_a_failed_docker_cp_raises_with_its_stderr(): with pytest.raises(RuntimeError, match="(?s)exit 1.*no such container"): await _ArgsRecorder(exit_code=1).docker_cp("/tmp/a", "c:/x") + + +@pytest.mark.asyncio +async def test_the_s3_named_methods_are_deprecated_aliases(signing_store): + sandbox = _RecordingVmSandbox() + with pytest.warns(DeprecationWarning, match="use load_object_file"): + await sandbox.load_s3_file("s3://bucket/a.json", "/tmp/a.json") + with pytest.warns(DeprecationWarning, match="use write_file_from_object"): + await sandbox.write_file_from_s3("s3://bucket/b.json", "/work/b.json") + assert any("https://signed/a.json" in s for s in sandbox.scripts) + assert any("https://signed/b.json" in s for s in sandbox.scripts) + + +@pytest.mark.asyncio +async def test_a_url_the_local_store_holds_is_read_from_it_not_the_configured_store(cli_routing): + """Under namespace routing, outside an @local run, the configured store owns only its own urls.""" + set_object_store(ConfiguredObjectStore()) + local = get_config().get_object_store_for("@local/~/t") + data = local.put("objects/data.json", b"LOCAL-DATA") + image = local.put("images/img.tar.gz", b"LOCAL-IMAGE") + sandbox = _RecordingVmSandbox(images_stdout="myimage\n") + + await sandbox.load_object_file(data, "/tmp/data.json") + await sandbox.load_docker_images([SimpleNamespace(tar_gz_object_url=image, image_name="myimage:latest")]) + + pushed = "\n".join(sandbox.scripts) + assert base64.b64encode(b"LOCAL-DATA").decode() in pushed + assert base64.b64encode(b"LOCAL-IMAGE").decode() in pushed diff --git a/tst/unit/store/fakes.py b/tst/unit/store/fakes.py index 4803e56c..4322fd39 100644 --- a/tst/unit/store/fakes.py +++ b/tst/unit/store/fakes.py @@ -157,6 +157,27 @@ def get_object_key(self, object_url): return object_url[len(home):] +class ConfiguredObjectStore(FakeObjectStore): + """The configured store beside the local one under namespace routing. It fails on any url it + doesn't own, so a test can prove a handed-in url was read or signed through the store holding it.""" + + def _own(self, object_url): + if not self.owns(object_url): + raise AssertionError(f"the configured store was asked about {object_url}") + + def signed_get_url(self, object_url, expires_in=3600): + self._own(object_url) + return super().signed_get_url(object_url, expires_in) + + def get(self, object_url): + self._own(object_url) + return super().get(object_url) + + def open(self, object_url): + self._own(object_url) + return super().open(object_url) + + class RecordingObjectStore(FakeObjectStore): """A FakeObjectStore that records the keys each key-based read addressed. diff --git a/tst/unit/task_step/run_docker_container_test.py b/tst/unit/task_step/run_docker_container_test.py index a77014e6..ba92180c 100644 --- a/tst/unit/task_step/run_docker_container_test.py +++ b/tst/unit/task_step/run_docker_container_test.py @@ -133,7 +133,7 @@ def __init__(self): self.loaded: list[tuple[str, str]] = [] self.scripts: list[str] = [] - async def load_s3_file(self, url, destination_path): + async def load_object_file(self, url, destination_path): self.loaded.append((url, destination_path)) async def exec_script(self, script): diff --git a/tst/unit/task_step/test_load_artifact_container_sandbox.py b/tst/unit/task_step/test_load_artifact_container_sandbox.py index 5d80e64f..24e68243 100644 --- a/tst/unit/task_step/test_load_artifact_container_sandbox.py +++ b/tst/unit/task_step/test_load_artifact_container_sandbox.py @@ -68,12 +68,12 @@ async def test_a_container_sandbox_is_written_directly(universe, monkeypatch): assert sorted(c.args for c in sandbox.exec.await_args_list) == [ ("mkdir", "-p", "/work/inputs"), ("mkdir", "-p", "/work/inputs/kb"), ] - assert sorted(c.args for c in sandbox.write_file_from_s3.await_args_list) == [ + assert sorted(c.args for c in sandbox.write_file_from_object.await_args_list) == [ ("s3://bucket/INDEX.md", "/work/inputs/kb/INDEX.md"), ("s3://bucket/task.json", "/work/inputs/task.json"), ] sandbox.exec_script.assert_not_awaited() - sandbox.load_s3_file.assert_not_awaited() + sandbox.load_object_file.assert_not_awaited() (loaded,) = ctx.metadata["loaded_file_artifact_universes"] assert loaded["sandbox_name"] == "bundler" assert loaded["destination_path"] == "/work/inputs" @@ -91,8 +91,8 @@ async def test_a_container_on_a_vm_backed_provider_is_loaded_through_its_contain scripts = [c.args[0] for c in sandbox.exec_script.await_args_list] assert scripts and all(s.startswith("docker exec -u 0 agent-local-1 mkdir -p ") for s in scripts) - assert len(sandbox.write_file_from_s3.await_args_list) == 2 - sandbox.load_s3_file.assert_not_awaited() + assert len(sandbox.write_file_from_object.await_args_list) == 2 + sandbox.load_object_file.assert_not_awaited() @pytest.mark.asyncio diff --git a/tst/unit/task_step/test_load_artifact_file.py b/tst/unit/task_step/test_load_artifact_file.py index 83fe8a63..35b581a6 100644 --- a/tst/unit/task_step/test_load_artifact_file.py +++ b/tst/unit/task_step/test_load_artifact_file.py @@ -39,11 +39,11 @@ async def exec(self, *command): async def docker_cp(self, source, destination, *, remove_source=False): self.commands.append(f"docker cp {source} {destination}") - async def load_s3_file(self, s3_url, destination_path): - self.pulls.append((s3_url, destination_path)) + async def load_object_file(self, object_url, destination_path): + self.pulls.append((object_url, destination_path)) - async def write_file_from_s3(self, s3_url, destination_path): - self.pulls.append((s3_url, destination_path)) + async def write_file_from_object(self, object_url, destination_path): + self.pulls.append((object_url, destination_path)) def _provider(sandbox: _Sandbox): diff --git a/tst/unit/task_step/test_load_artifact_vm_target.py b/tst/unit/task_step/test_load_artifact_vm_target.py index c69da642..30ec4d85 100644 --- a/tst/unit/task_step/test_load_artifact_vm_target.py +++ b/tst/unit/task_step/test_load_artifact_vm_target.py @@ -41,8 +41,8 @@ async def exec_script(self, script, **kw): calls.append(("exec", script)) return "" - async def load_s3_file(self, s3_url, destination_path): - calls.append(("s3", s3_url, destination_path)) + async def load_object_file(self, object_url, destination_path): + calls.append(("object", object_url, destination_path)) async def load_url_file(self, url, destination_path): # pragma: no cover # Deliberately present: the step must NOT depend on a sandbox method. @@ -101,11 +101,11 @@ async def test_files_land_on_the_vm_with_no_docker(self, universe, vm): ctx = await step.execute(ctx) assert not any("docker" in c[1] for c in calls if c[0] == "exec"), calls - # Pulled straight from S3 to the final path — no VM temp file, no copy inward. - s3 = [c for c in calls if c[0] == "s3"] - assert s3 == [ - ("s3", "s3://bucket/instruction.md", "/app/bundle/instruction.md"), - ("s3", "s3://bucket/test.sh", "/app/bundle/tests/test.sh"), + # Pulled straight from the object store to the final path — no VM temp file, no copy inward. + objects = [c for c in calls if c[0] == "object"] + assert objects == [ + ("object", "s3://bucket/instruction.md", "/app/bundle/instruction.md"), + ("object", "s3://bucket/test.sh", "/app/bundle/tests/test.sh"), ] assert ("exec", "mkdir -p /app/bundle") in calls assert ("exec", "mkdir -p /app/bundle/tests") in calls diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index ca96aab2..99d51981 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -10,11 +10,13 @@ import agent_env.config as config_mod from agent_env.env import legacy_protocol from agent_env.env.env import DeployedGatewayEnv +from agent_env.store import set_object_store from agent_env.store.object_store import S3ObjectStore from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps import snapshot_agent_state as mod from agentenv_protocol import client as protocol_v1 from tst.unit.event_loop_probe import on_event_loop +from tst.unit.store.fakes import ConfiguredObjectStore _PREFIX = "s3://artifact-bucket/agent_snapshots/oc_post_run_workspace_T/3-deadbeef/" @@ -87,7 +89,7 @@ def get(_id): import agent_env.config as config_mod class _Cfg: - def get_object_store(self): + def get_object_store_at(self, object_url): return store monkeypatch.setattr(config_mod, "get_config", lambda: _Cfg()) @@ -125,7 +127,7 @@ def put_object(self, **kw): store = S3ObjectStore(_LoopCheckingS3(), "artifact-bucket") monkeypatch.setattr(artifact_mod, "EnvironmentUniverseArtifact", type("U", (), {"get": staticmethod(lambda _id: _StubUniverse())})) - monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store": lambda self: store})()) + monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store_at": lambda self, object_url: store})()) async def _no_v1(_base): return False @@ -151,7 +153,7 @@ async def test_capture_reads_child_envs_on_the_stored_card_and_takes_legacy_for_ store = S3ObjectStore(_StubS3(), "artifact-bucket") monkeypatch.setattr(artifact_mod, "EnvironmentUniverseArtifact", type("U", (), {"get": staticmethod(lambda _id: _StubUniverse())})) - monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store": lambda self: store})()) + monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store_at": lambda self, object_url: store})()) async def _no_probe(_base): raise AssertionError("a stored card means no live probe") @@ -195,7 +197,7 @@ def put_object(self, Bucket, Key, Body, ContentType): # noqa: N803 store = S3ObjectStore(_BodyS3(), "artifact-bucket") monkeypatch.setattr(artifact_mod, "EnvironmentUniverseArtifact", type("U", (), {"get": staticmethod(lambda _id: _StubUniverse())})) - monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store": lambda self: store})()) + monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store_at": lambda self, object_url: store})()) async def _get_data(base, timeout=30): if base.endswith("/mcp-calendar"): @@ -241,3 +243,24 @@ def _carded_ctx(card: dict) -> TaskStepContext: ctx.deployed_envs[0].environment_card_url = f"https://sandbox.example/sb-1{WELL_KNOWN_PATH}" ctx.deployed_envs[0].environment_card = card return ctx + + +@pytest.mark.asyncio +async def test_services_land_beside_a_capture_the_local_store_holds(monkeypatch, cli_routing): + set_object_store(ConfiguredObjectStore()) + local = config_mod.get_config().get_object_store_for("@local/~/t") + prefix = local.object_url("agent_snapshots/oc/3-deadbeef/") + monkeypatch.setattr(artifact_mod, "EnvironmentUniverseArtifact", type("U", (), {"get": staticmethod(lambda _id: _StubUniverse())})) + + async def _state(_deployed, _gateway, name): + return {"service": name} + + monkeypatch.setattr(legacy_protocol, "service_state", _state) + ctx = _ctx() + await _step()._capture_universe_state(ctx, prefix) + + published = json.loads(ctx.metadata["snapshot_json_url"]) + assert published == { + name: local.object_url(f"agent_snapshots/oc/3-deadbeef/services/{name}.json") for name in ("calendar", "contacts") + } + assert json.loads(local.get(published["calendar"])) == {"service": "calendar"} diff --git a/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py b/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py index 8c8614a2..511a120f 100644 --- a/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py +++ b/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py @@ -24,9 +24,10 @@ from agent_env.task_step.context import PromptResponse, TaskStepContext from agent_env.task_step.task_steps.verifiers import rubrics_verifier from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep -from agent_env.config import set_object_store +from agent_env.config import get_config, set_object_store from agent_env.task_step.task_steps.verifiers.judge_utils.trajectory_filter import CompactionType, TrajectoryFilter from tst.unit.event_loop_probe import on_event_loop +from tst.unit.store.fakes import ConfiguredObjectStore def _verifier(**kw) -> RubricsVerifierTaskStep: @@ -187,13 +188,13 @@ def _no_merge(*a, **k): async def test_load_per_turn_trajectories_no_filter_writes_each_turn(): v = _verifier() sandbox = MagicMock() - sandbox.write_file_from_s3 = AsyncMock() + sandbox.write_file_from_object = AsyncMock() sandbox.write_file_from_text = AsyncMock() await v._load_per_turn_trajectories( sandbox, ["s3://a/1.json", "s3://a/2.json"], "/tmp/d", None) - dests = [call.args[1] for call in sandbox.write_file_from_s3.call_args_list] + dests = [call.args[1] for call in sandbox.write_file_from_object.call_args_list] assert dests == ["/tmp/d/turn_01.json", "/tmp/d/turn_02.json"] sandbox.write_file_from_text.assert_not_called() @@ -204,12 +205,12 @@ async def test_load_per_turn_trajectories_default_filter_compacts_and_namespaces v = _verifier() monkeypatch.setattr(v, "_read_trajectory_text", lambda uri: spans) sandbox = MagicMock() - sandbox.write_file_from_s3 = AsyncMock() + sandbox.write_file_from_object = AsyncMock() sandbox.write_file_from_text = AsyncMock() await v._load_per_turn_trajectories(sandbox, ["u1", "u2"], "/tmp/d", TrajectoryFilter()) - sandbox.write_file_from_s3.assert_not_called() # compacted -> written as text + sandbox.write_file_from_object.assert_not_called() # compacted -> written as text dests = [call.args[1] for call in sandbox.write_file_from_text.call_args_list] assert "/tmp/d/turn_01.json" in dests and "/tmp/d/turn_02.json" in dests tool_dests = [d for d in dests if "tool_call_result" in d] @@ -355,3 +356,17 @@ async def judge(*, eval_prompt, **kwargs): await v.execute(TaskStepContext(prompt_responses=[_pr(agent_trajectory_s3_uri="s3://t/raw.json")])) assert held and alive == [False] + + +def test_a_trajectory_the_local_store_holds_is_read_from_it(cli_routing): + """Read from the store holding the handed-in url; the compacted copy is a new object, minted in the + configured store.""" + configured = ConfiguredObjectStore() + set_object_store(configured) + spans = [_tool_span("t", {"result": "x" * 2000})] + url = get_config().get_object_store_for("@local/~/t").put("trajectories/raw.json", json.dumps(spans).encode()) + v = _verifier(trajectory_filter=TrajectoryFilter()) + + assert json.loads(v._read_trajectory_text(url)) == spans + compacted, _ = v._filter_trajectory(url, TrajectoryFilter()) + assert configured.owns(compacted) diff --git a/tst/unit/utils/test_paths.py b/tst/unit/utils/test_paths.py index 4517116c..1f82c7eb 100644 --- a/tst/unit/utils/test_paths.py +++ b/tst/unit/utils/test_paths.py @@ -41,7 +41,7 @@ class _Sandbox: async def exec(self, *a): raise AssertionError("must not reach the sandbox") - async def write_file_from_s3(self, *a): + async def write_file_from_object(self, *a): raise AssertionError("must not write anything") class _Provider: