From ed59a6ec2a7028b0ba3491317e7aec0252489f57 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 13:31:46 -0700 Subject: [PATCH 01/11] feat(snapshot_env): hand a service that takes parts grants an object of its own to upload to When a service's card lists http-put-parts (client.write_kinds), snapshot_env begins a write of agentenv-snapshots//.zip (application/zip, up to 16 GiB) and sends it as data/get's write_object. On an uploaded_object reply it completes the write and registers the bundle in place, with the same key shape and file name as a self-upload; any other reply, an error or a timeout aborts it. A completion failure is an UploadFailedError, a ValueError, so it fails only that service. Services that do not declare parts are offered what they are today. A service that also advertises the credentials extension still gets the pushed credentials, so a parts upload that fails can fall back to them. Each export now logs one line naming how the bundle came (parts, namespace, service upload, inline, stream, data or export-state) and how long it took, with the bytes and parts of a parts upload. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../task_step/task_steps/snapshot_env.py | 99 +++++++-- .../snapshot_env_object_write_test.py | 193 ++++++++++++++++++ 2 files changed, 272 insertions(+), 20 deletions(-) create mode 100644 tst/unit/task_step/snapshot_env_object_write_test.py diff --git a/src/agent_env/task_step/task_steps/snapshot_env.py b/src/agent_env/task_step/task_steps/snapshot_env.py index c78c5042f..e5256f5e3 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -5,9 +5,11 @@ import base64 import json import logging +import math import mimetypes import os import tempfile +import time import uuid from dataclasses import dataclass from datetime import datetime, timezone @@ -16,8 +18,8 @@ import httpx -from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object_path -from agentenv_protocol.client import find_extension +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object, uploaded_object_path +from agentenv_protocol.client import find_extension, write_kinds from agentenv_protocol.transfers import WriteNamespaceGrant from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant @@ -25,7 +27,7 @@ from agent_env.store import get_config from agent_env.store.base import GrantUnavailableError from agent_env.store.ids import derive_id, is_local_id, validate_local_id -from agent_env.store.object_store import MIN_GRANT_LIFETIME_SECONDS, ObjectStore +from agent_env.store.object_store import MIN_GRANT_LIFETIME_SECONDS, ObjectStore, PendingWrite from agent_env.store.object_store.object_store import issues_grants_to from agent_env.store.routing import in_local_run from agent_env.task_step.context import TaskStepContext @@ -46,6 +48,9 @@ # What a service may upload through its snapshot grant: one bundle, as large as one upload to the # store can be up to this. ENV_SNAPSHOT_LIMITS = ObjectLimits(max_objects=1, max_object_bytes=50 * 1024**3, max_total_bytes=50 * 1024**3) +# The most a service may upload as its bundle through an object write: room for the largest bundles seen, +# about twice over, in a grant of a few hundred part URLs. +ENV_SNAPSHOT_MAX_BYTES = 16 * 1024**3 _SNAPSHOT_KEY_PREFIX = "agentenv-snapshots" @@ -87,6 +92,35 @@ def _snapshot_upload(card: dict, timeout_seconds: float, sandbox_type: Optional[ return _SnapshotUpload(store, grant) +def _snapshot_write( + card: dict, environment_name: str, timeout_seconds: float, sandbox_type: Optional[str] +) -> Optional[PendingWrite]: + """A write of a service's bundle to an object of its own, for a service whose card takes parts grants; + None when it takes none, or the object store issues no grant that reaches the service's sandbox, of + ``sandbox_type`` (None: unknown), and lasts the export.""" + kinds = write_kinds(card) + if "http-put-parts" not in kinds: + return None + config = get_config() + store = config.get_object_store() + if not issues_grants_to(store, sandbox_type): + return None + object_url = store.object_url( + f"{config.get_artifact_key_prefix()}{_SNAPSHOT_KEY_PREFIX}/{uuid.uuid4().hex}/{environment_name}.zip" + ) + try: + return store.begin_write( + object_url, + media_type="application/zip", + max_bytes=ENV_SNAPSHOT_MAX_BYTES, + kinds=kinds, + expires_in=max(int(timeout_seconds), MIN_GRANT_LIFETIME_SECONDS), + ) + except GrantUnavailableError as exc: + logger.info("snapshot_env: the object store issues no snapshot write grant (%s)", exc) + return None + + async def _push_s3_credentials(base_url: str, card: dict, timeout_seconds: float) -> None: """Best-effort: push the object store's shared AWS creds + bucket to a service advertising add_s3_credentials, if the store shares any. Never raises (get_data falls back). @@ -412,28 +446,52 @@ async def _export_environment_to_file( ``tmp_path``, or the object url a FilePart names when the service uploaded the bundle itself (nothing written to ``tmp_path`` then — export_one registers it). - v1 ``get_data``, handed a namespace grant when the service takes the data-objects - form: an upload through it → its object url; ``DataPart`` → json; ``FilePart`` → - object url (no download) / base64 ``bytes`` / streamed relative-or-http ``uri`` → - ".zip"; else legacy ``GET /export-state`` streamed to disk with a first-byte JSON - guard. Errors are caught by the caller (fails just this service).""" + v1 ``get_data``, handed an object write when the service takes parts grants, else a + namespace grant when it takes the data-objects form: an upload through either → its + object url; ``DataPart`` → json; ``FilePart`` → object url (no download) / base64 + ``bytes`` / streamed relative-or-http ``uri`` → ".zip"; else legacy + ``GET /export-state`` streamed to disk with a first-byte JSON guard. Logs which way the + export came and how long it took. Errors are caught by the caller (fails just this + service).""" from agent_env.env import legacy_protocol from agentenv_protocol import client as protocol_v1 + started = time.monotonic() + + def exported(tier: str, result: str, detail: str = "") -> str: + logger.info( + "snapshot_env: %s exported through %s%s in %.1fs", + environment_name, tier, detail, time.monotonic() - started, + ) + return result + base_url = await legacy_protocol.v1_base_url(deployed, gateway_url, environment_name, mcp=True) if base_url is not None: card_base, card = await legacy_protocol.child_env_card(deployed, gateway_url, environment_name) - upload = await asyncio.to_thread( - _snapshot_upload, card or {}, timeout_seconds, getattr(deployed, "sandbox_type", None) + sandbox_type = getattr(deployed, "sandbox_type", None) + write = await asyncio.to_thread(_snapshot_write, card or {}, environment_name, timeout_seconds, sandbox_type) + upload = None if write is not None else await asyncio.to_thread( + _snapshot_upload, card or {}, timeout_seconds, sandbox_type ) - # Push S3 creds (no-op unless the service advertises the extension), with a grant too: a - # bundle the grant cannot hold can still be uploaded with them. + # Push S3 creds (no-op unless the service advertises the extension) beside either grant: a + # bundle the grant cannot hold, or an upload through it that fails, can still use them. await _push_s3_credentials(card_base, card or {}, timeout_seconds) - grant = {} if upload is None else {"write_namespace": upload.grant} - resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) - part = resp.parts[0] if resp.parts else None + if write is not None: + grant = {"write_object": write.grant} + else: + grant = {} if upload is None else {"write_namespace": upload.grant} + try: + resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) + part = resp.parts[0] if resp.parts else None + if write is not None and (uploaded := uploaded_object(part)) is not None: + await asyncio.to_thread(write.complete, uploaded) + parts = max(1, math.ceil(uploaded.size_bytes / write.grant.write.part_bytes)) + return exported("parts", write.object_url, f" ({uploaded.size_bytes} bytes in {parts} parts)") + finally: + if write is not None: + await asyncio.to_thread(write.abort) if upload is not None and (path := uploaded_object_path(part)) is not None: - return upload.object_url(path) + return exported("namespace", upload.object_url(path)) if isinstance(part, FilePart): raw = getattr(part.file, "bytes", None) if raw is not None: @@ -452,7 +510,7 @@ async def _export_environment_to_file( # last segment carries the shape (e.g. gdrive.zip). scheme = urlparse(uri).scheme if scheme and scheme not in ("http", "https"): - return uri + return exported("service upload", uri) # A relative uri (e.g. "export-snapshot") means "stream it from # my service endpoint" — resolve against the service base_url so # multi-GB bundles never ride inline in the JSON-RPC response. @@ -470,11 +528,12 @@ async def _export_environment_to_file( with open(tmp_path, "wb") as f: async for chunk in resp_file.aiter_bytes(): f.write(chunk) - return os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" + suffix = os.path.splitext(getattr(part.file, "name", "") or "")[1] or ".zip" + return exported("inline" if raw is not None else "stream", suffix) state = part.data if part is not None else {} with open(tmp_path, "w") as f: json.dump(state, f, default=str) - return ".json" + return exported("data", ".json") legacy_url = f"{legacy_protocol.environment_base_url(gateway_url, environment_name, mcp=True)}/export-state" async with httpx.AsyncClient() as client: @@ -495,7 +554,7 @@ async def _export_environment_to_file( ) seen_first_bytes = True f.write(chunk) - return ".json" + return exported("export-state", ".json") @classmethod async def snapshot_env_state( diff --git a/tst/unit/task_step/snapshot_env_object_write_test.py b/tst/unit/task_step/snapshot_env_object_write_test.py new file mode 100644 index 000000000..6c9d5bd29 --- /dev/null +++ b/tst/unit/task_step/snapshot_env_object_write_test.py @@ -0,0 +1,193 @@ +"""snapshot_env hands a service that takes parts grants an object of its own to upload its bundle to, makes +the object from what it reports, and registers it there; every other reply aborts the write.""" + +from __future__ import annotations + +import dataclasses +import json +import logging +from urllib.parse import unquote, urlsplit + +import httpx +import pytest +from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, TRANSFERS_EXTENSION_URI, uploaded_object_part +from agentenv_protocol.transfers import Uploaded + +from agent_env.artifact import EnvironmentUniverseArtifact +from agent_env.config import set_object_store +from agent_env.store.base import GrantUnavailableError +from agent_env.task_step.task_steps import snapshot_env as snapshot_mod +from agent_env.task_step.task_steps.snapshot_env import ENV_SNAPSHOT_MAX_BYTES, SnapshotEnvTaskStep +from tst.unit.task_step.test_snapshot_env import _CARD, _FakeMultiEnv, _card, _carded +from tst.util.granting_object_store import GrantingObjectStore + +_PARTS = {"uri": TRANSFERS_EXTENSION_URI, "params": {"write": ["http-post-policy", "http-put", "http-put-parts"]}} +_DATA_OBJECTS = {"uri": DATA_OBJECTS_EXTENSION_URI} +_S3_CREDENTIALS = {"uri": "urn:agentenv:add-s3-credentials/v1", + "params": {"endpoint": "/agentenv/ext/add_s3_credentials", "methods": {"add_s3_credentials": {"method": "POST"}}}} +_BUNDLE = b"PK\x03\x04bundle" + + +class _UngrantingStore(GrantingObjectStore): + def begin_write(self, object_url, **kwargs): + raise GrantUnavailableError("this store issues no object writes") + + +class _WatchedStore(GrantingObjectStore): + """Records how each write it begins ends.""" + + def __init__(self, root): + super().__init__(root) + self.ended: list[str] = [] + + def begin_write(self, object_url, **kwargs): + write = super().begin_write(object_url, **kwargs) + complete, abort, ended = write.complete, write.abort, self.ended + write.complete = lambda uploaded: (complete(uploaded), ended.append("completed"))[0] + write.abort = lambda: (ended.append("aborted") if not write._settled else None, abort())[1] + return write + + +@pytest.fixture +def store(local_stores, tmp_path): + granting = _WatchedStore(str(tmp_path / "objects")) + set_object_store(granting) + return granting + + +@pytest.fixture +def pushes(monkeypatch): + pushed: list[str] = [] + + async def push(base_url, card, timeout_seconds): + pushed.append(base_url) + + monkeypatch.setattr(snapshot_mod, "_push_s3_credentials", push) + return pushed + + +async def _record(children: dict, *, sandbox_type: str | None = "local"): + return dataclasses.replace(_carded(await _card(children)), sandbox_type=sandbox_type) + + +def _serve(monkeypatch, answers: dict): + """Answer each service's data/get with ``answers[service](params)``, recording the params sent.""" + sent, real = [], httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + service = next(seg for seg in request.url.path.split("/") if seg.startswith("mcp-")).removeprefix("mcp-") + params = json.loads(request.content)["params"] + sent.append((service, params)) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": answers[service](params)}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + return sent + + +def _key(params) -> str: + return unquote(urlsplit(params["write_object"]["write"]["urls"][0]).path.lstrip("/")) + + +def _uploading(store, *, reported: int | None = None): + """A service that PUTs its bundle to the object it is handed and says what it sent.""" + def answer(params): + store.put(_key(params), _BUNDLE, content_type="application/zip") + uploaded = Uploaded(size_bytes=len(_BUNDLE) if reported is None else reported) + part = uploaded_object_part(uploaded, name="gdrive.zip", mime_type="application/zip") + return {"parts": [part.model_dump(mode="json", exclude_none=True)]} + return answer + + +def _data(params): + return {"parts": [{"kind": "data", "data": {"rows": 1}}]} + + +@pytest.mark.asyncio +async def test_a_service_taking_parts_uploads_to_an_object_of_its_own_and_is_registered_there( + store, pushes, monkeypatch, caplog +): + record = await _record({"mcp-gdrive": [_PARTS, _DATA_OBJECTS, _S3_CREDENTIALS]}) + sent = _serve(monkeypatch, {"gdrive": _uploading(store)}) + + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + result = await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive"]), gateway_url="https://gw-1", snapshot_id="snap-o", deployed=record, + ) + + ((_, params),) = sent + assert list(params) == ["write_object"] + grant = params["write_object"] + assert (grant["media_type"], grant["max_bytes"], grant["write"]["kind"]) == ("application/zip", ENV_SNAPSHOT_MAX_BYTES, "http-put-parts") + key = _key(params) + assert key.startswith("agentenv-snapshots/") and key.endswith("/gdrive.zip") + assert pushes == [_CARD] + assert store.ended == ["completed"] + (bundle,) = EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id).get_file_artifacts().values() + assert bundle.object_url == store.object_url(key) + assert (bundle.filename, bundle.content_type, bundle.load()) == ("gdrive.zip", "application/zip", _BUNDLE) + assert f"gdrive exported through parts ({len(_BUNDLE)} bytes in 1 parts)" in caplog.text + + +@pytest.mark.asyncio +async def test_each_object_write_gets_an_object_of_its_own(store, pushes, monkeypatch): + record = await _record({"mcp-gdrive": [_PARTS]}) + sent = _serve(monkeypatch, {"gdrive": _uploading(store)}) + + for _ in range(2): + await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive"]), gateway_url="https://gw-1", snapshot_id="snap-1", deployed=record, + ) + + first, second = (_key(params) for _, params in sent) + assert first != second + + +@pytest.mark.asyncio +async def test_a_reply_that_names_no_object_upload_aborts_the_write_and_exports_as_before( + store, pushes, monkeypatch, tmp_path +): + record = await _record({"mcp-gdrive": [_PARTS]}) + _serve(monkeypatch, {"gdrive": _data}) + + out = tmp_path / "export-tmp" + suffix = await SnapshotEnvTaskStep._export_environment_to_file("https://gw-1", "gdrive", str(out), 30, record) + + assert (suffix, json.loads(out.read_text())) == (".json", {"rows": 1}) + assert store.ended == ["aborted"] + + +@pytest.mark.asyncio +async def test_an_upload_the_store_cannot_make_into_the_object_fails_only_that_service(store, pushes, monkeypatch): + record = await _record({"mcp-gdrive": [_PARTS], "mcp-slack": []}) + _serve(monkeypatch, {"gdrive": _uploading(store, reported=len(_BUNDLE) + 1), "slack": _data}) + + with pytest.raises(RuntimeError, match=r"1/2 services were not exported.*gdrive"): + await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive", "slack"]), gateway_url="https://gw-1", snapshot_id="snap-f", deployed=record, + ) + assert store.ended == ["aborted"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("extensions", "store_cls", "sandbox_type", "params"), [ + ([_DATA_OBJECTS], _WatchedStore, "local", "write_namespace"), + ([_PARTS, _DATA_OBJECTS], _WatchedStore, "modal", None), + ([_PARTS, _DATA_OBJECTS], _UngrantingStore, "local", "write_namespace"), +], ids=["parts-not-declared", "grants-do-not-reach-the-sandbox", "store-issues-no-object-write"]) +async def test_otherwise_the_export_is_offered_as_before( + local_stores, tmp_path, pushes, monkeypatch, caplog, extensions, store_cls, sandbox_type, params +): + set_object_store(store_cls(str(tmp_path / "objects"))) + record = await _record({"mcp-gdrive": extensions}, sandbox_type=sandbox_type) + sent = _serve(monkeypatch, {"gdrive": _data}) + + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + suffix = await SnapshotEnvTaskStep._export_environment_to_file( + "https://gw-1", "gdrive", str(tmp_path / "export-tmp"), 30, record + ) + + assert suffix == ".json" + ((_, sent_params),) = sent + assert list(sent_params) == ([] if params is None else [params]) + assert pushes == [_CARD] + assert "gdrive exported through data in" in caplog.text From 6c052d24d848c0c440a7eca6814d2c2c9d3da0eb Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 13:54:16 -0700 Subject: [PATCH 02/11] fix(snapshot): an export cancelled midway leaves no write open and no abort racing its completion A write begun or completing in a thread when the export is cancelled is now seen through: the export waits for the thread, aborts a write that began after the cancellation, and lets a completion finish before its cleanup runs. The creds push moved inside the try, so a cancellation there aborts the write. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../task_step/task_steps/snapshot_env.py | 38 +++++++++++----- .../snapshot_env_object_write_test.py | 45 +++++++++++++++++++ 2 files changed, 71 insertions(+), 12 deletions(-) diff --git a/src/agent_env/task_step/task_steps/snapshot_env.py b/src/agent_env/task_step/task_steps/snapshot_env.py index e5256f5e3..18db4564b 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -121,6 +121,19 @@ def _snapshot_write( return None +async def _finish_in_thread(fn, *args): + """``await asyncio.to_thread(fn, *args)``, except that a caller cancelled meanwhile first waits for ``fn`` to + return, and aborts a write it began, so neither a write nor its completion outlives the caller's cleanup.""" + task = asyncio.ensure_future(asyncio.to_thread(fn, *args)) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + await asyncio.wait([task]) + if task.exception() is None and isinstance(task.result(), PendingWrite): + await asyncio.to_thread(task.result().abort) + raise + + async def _push_s3_credentials(base_url: str, card: dict, timeout_seconds: float) -> None: """Best-effort: push the object store's shared AWS creds + bucket to a service advertising add_s3_credentials, if the store shares any. Never raises (get_data falls back). @@ -469,22 +482,23 @@ def exported(tier: str, result: str, detail: str = "") -> str: if base_url is not None: card_base, card = await legacy_protocol.child_env_card(deployed, gateway_url, environment_name) sandbox_type = getattr(deployed, "sandbox_type", None) - write = await asyncio.to_thread(_snapshot_write, card or {}, environment_name, timeout_seconds, sandbox_type) - upload = None if write is not None else await asyncio.to_thread( - _snapshot_upload, card or {}, timeout_seconds, sandbox_type - ) - # Push S3 creds (no-op unless the service advertises the extension) beside either grant: a - # bundle the grant cannot hold, or an upload through it that fails, can still use them. - await _push_s3_credentials(card_base, card or {}, timeout_seconds) - if write is not None: - grant = {"write_object": write.grant} - else: - grant = {} if upload is None else {"write_namespace": upload.grant} + write = None try: + write = await _finish_in_thread(_snapshot_write, card or {}, environment_name, timeout_seconds, sandbox_type) + upload = None if write is not None else await asyncio.to_thread( + _snapshot_upload, card or {}, timeout_seconds, sandbox_type + ) + # Push S3 creds (no-op unless the service advertises the extension) beside either grant: a + # bundle the grant cannot hold, or an upload through it that fails, can still use them. + await _push_s3_credentials(card_base, card or {}, timeout_seconds) + if write is not None: + grant = {"write_object": write.grant} + else: + grant = {} if upload is None else {"write_namespace": upload.grant} resp = await protocol_v1.get_data(base_url, timeout=int(timeout_seconds), **grant) part = resp.parts[0] if resp.parts else None if write is not None and (uploaded := uploaded_object(part)) is not None: - await asyncio.to_thread(write.complete, uploaded) + await _finish_in_thread(write.complete, uploaded) parts = max(1, math.ceil(uploaded.size_bytes / write.grant.write.part_bytes)) return exported("parts", write.object_url, f" ({uploaded.size_bytes} bytes in {parts} parts)") finally: diff --git a/tst/unit/task_step/snapshot_env_object_write_test.py b/tst/unit/task_step/snapshot_env_object_write_test.py index 6c9d5bd29..adbf93b0d 100644 --- a/tst/unit/task_step/snapshot_env_object_write_test.py +++ b/tst/unit/task_step/snapshot_env_object_write_test.py @@ -3,9 +3,11 @@ from __future__ import annotations +import asyncio import dataclasses import json import logging +import threading from urllib.parse import unquote, urlsplit import httpx @@ -48,6 +50,26 @@ def begin_write(self, object_url, **kwargs): return write +class _GatedStore(_WatchedStore): + """Holds the export at ``at`` (beginning the write, pushing creds, completing it) until released.""" + + def __init__(self, root, at): + super().__init__(root) + self.at, self.reached, self.release = at, threading.Event(), threading.Event() + + def hold(self, at): + if at == self.at: + self.reached.set() + self.release.wait(5) + + def begin_write(self, object_url, **kwargs): + self.hold("begin") + write = super().begin_write(object_url, **kwargs) + complete = write.complete + write.complete = lambda uploaded: (self.hold("complete"), complete(uploaded))[1] + return write + + @pytest.fixture def store(local_stores, tmp_path): granting = _WatchedStore(str(tmp_path / "objects")) @@ -191,3 +213,26 @@ async def test_otherwise_the_export_is_offered_as_before( assert list(sent_params) == ([] if params is None else [params]) assert pushes == [_CARD] assert "gdrive exported through data in" in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("at", "ended"), [("begin", ["aborted"]), ("push", ["aborted"]), ("complete", ["completed"])]) +async def test_an_export_cancelled_midway_leaves_no_write_open_or_aborted_mid_completion( + local_stores, tmp_path, monkeypatch, at, ended +): + store = _GatedStore(str(tmp_path / "objects"), at) + set_object_store(store) + monkeypatch.setattr(snapshot_mod, "_push_s3_credentials", lambda *_: asyncio.to_thread(store.hold, "push")) + record = await _record({"mcp-gdrive": [_PARTS]}) + _serve(monkeypatch, {"gdrive": _uploading(store)}) + + export = asyncio.ensure_future( + SnapshotEnvTaskStep._export_environment_to_file("https://gw-1", "gdrive", str(tmp_path / "out"), 30, record) + ) + assert await asyncio.to_thread(store.reached.wait, 5) + export.cancel() + await asyncio.sleep(0.05) + store.release.set() + with pytest.raises(asyncio.CancelledError): + await export + assert store.ended == ended From 53bcef6c63beb5bfdc136ed993f2b7f1aa77136b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 15:35:17 -0700 Subject: [PATCH 03/11] refactor(snapshot_env): count a parts upload's parts with part_ranges Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/task_step/task_steps/snapshot_env.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/agent_env/task_step/task_steps/snapshot_env.py b/src/agent_env/task_step/task_steps/snapshot_env.py index 18db4564b..a84605ad4 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -5,7 +5,6 @@ import base64 import json import logging -import math import mimetypes import os import tempfile @@ -20,7 +19,7 @@ from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object, uploaded_object_path from agentenv_protocol.client import find_extension, write_kinds -from agentenv_protocol.transfers import WriteNamespaceGrant +from agentenv_protocol.transfers import WriteNamespaceGrant, part_ranges from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant from agent_env.env.gateway.constants import EXT_TRAJECTORY_URI @@ -499,7 +498,7 @@ def exported(tier: str, result: str, detail: str = "") -> str: part = resp.parts[0] if resp.parts else None if write is not None and (uploaded := uploaded_object(part)) is not None: await _finish_in_thread(write.complete, uploaded) - parts = max(1, math.ceil(uploaded.size_bytes / write.grant.write.part_bytes)) + parts = len(part_ranges(uploaded.size_bytes, write.grant.write.part_bytes)) return exported("parts", write.object_url, f" ({uploaded.size_bytes} bytes in {parts} parts)") finally: if write is not None: From 32b9ddfe2dbba5b26b930028f68eccf1d7a446c4 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 15:38:36 -0700 Subject: [PATCH 04/11] fix(local): services on a local compose stack trust the local transfer CA; snapshot e2e through a parts grant A MultiEnv on the local backend runs its MCP servers from a compose file that never gave them the local grant server's CA, so every grant handed to one failed TLS verification. Each MCP server now mounts the CA's trust files where LOCAL_TRUST_ENV points, as a single local container already gets them. The new slow-tier test snapshots a MultiEnv whose items server takes write_object, checks the bundle it uploaded is registered and holds its state, and loads the snapshot back. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../env_providers/env_gateway_provider.py | 9 +++ .../sandbox_providers/local_sandbox.py | 10 +-- .../env/envs/multi_env_snapshot_local_test.py | 63 +++++++++++++++++++ .../env_gateway_provider_test.py | 24 +++++++ 4 files changed, 101 insertions(+), 5 deletions(-) create mode 100644 tst/integration/env/envs/multi_env_snapshot_local_test.py diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 3e96fd46a..aceb4a10d 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -10,6 +10,7 @@ import time from dataclasses import dataclass, field from datetime import datetime, timezone +from pathlib import Path from typing import TYPE_CHECKING, Optional from agent_env.attribution import Attribution @@ -33,6 +34,7 @@ GATEWAY_SERVICE_NAME, ) from agent_env.providers.env_providers.env_provider import _SandboxEnvironmentProvider, _size_kwargs, _tool_names +from agent_env.providers.sandbox_providers.local_sandbox import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV, LocalSandbox, local_grant_trust from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, port_bindings from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SandboxProvider @@ -215,8 +217,11 @@ def create_docker_compose( host_ips: tuple[str, ...] = (), extra_hosts: tuple[str, ...] = (), mcp_server_name: str | None = None, + trust_dir: Path | None = None, ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. + With ``trust_dir``, the local transfer CA's trust files, each MCP server mounts them where + ``LOCAL_TRUST_ENV`` points, so its TLS clients trust the local grant server. Supports MCP servers, websites, or both. Website containers (backend + frontend) are added to the same docker network. The gateway receives WEBSITE_URLS env var for discovery. @@ -361,6 +366,9 @@ def create_docker_compose( if server.extra_env_vars: for key, value in server.extra_env_vars.items(): lines.append(f" - {key}={value}") + if trust_dir is not None: + lines.extend(f" - {key}={value}" for key, value in LOCAL_TRUST_ENV.items()) + lines.extend([" volumes:", f' - "{trust_dir}:{LOCAL_TRUST_DIR}:ro"']) lines.extend([ " healthcheck:", f' test: ["CMD-SHELL", "python3 -c \\"import socket; s=socket.create_connection((\'localhost\',{gateway_port}),2); s.close()\\""]', @@ -737,6 +745,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, + trust_dir=await asyncio.to_thread(local_grant_trust) if isinstance(sandbox, LocalSandbox) else None, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index 7289aa167..1a585075f 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -74,11 +74,11 @@ def _runs_in_container(cmd: list[str]) -> bool: # Where a container finds the local transfer CA's trust files, and the variables that point TLS clients at them: # SSL_CERT_FILE replaces a client's roots, so it gets the public roots plus the CA; NODE_EXTRA_CA_CERTS adds. -_TRUST_DIR = "/etc/agentenv" +LOCAL_TRUST_DIR = "/etc/agentenv" LOCAL_TRUST_ENV = { - "SSL_CERT_FILE": f"{_TRUST_DIR}/ca-bundle.pem", - "REQUESTS_CA_BUNDLE": f"{_TRUST_DIR}/ca-bundle.pem", - "NODE_EXTRA_CA_CERTS": f"{_TRUST_DIR}/ca.pem", + "SSL_CERT_FILE": f"{LOCAL_TRUST_DIR}/ca-bundle.pem", + "REQUESTS_CA_BUNDLE": f"{LOCAL_TRUST_DIR}/ca-bundle.pem", + "NODE_EXTRA_CA_CERTS": f"{LOCAL_TRUST_DIR}/ca.pem", } # Marker dropped in the work dir when this sandbox runs a container, so a later get_sandbox() @@ -460,7 +460,7 @@ async def start_trusting(sandbox: VmSandbox, container: str, trust_dir: Path) -> """Copy ``trust_dir`` into the created ``container`` where ``LOCAL_TRUST_ENV`` points, then start it. A container that cannot be given the files or started is removed, so its name is free for the next attempt.""" try: - await asyncio.to_thread(_copy_into_container, trust_dir, container, _TRUST_DIR) + await asyncio.to_thread(_copy_into_container, trust_dir, container, LOCAL_TRUST_DIR) await sandbox.exec_script(f"docker start {shlex.quote(container)} > /dev/null") except Exception: await sandbox.exec_script(f"docker rm -f {shlex.quote(container)} >/dev/null 2>&1 || true") diff --git a/tst/integration/env/envs/multi_env_snapshot_local_test.py b/tst/integration/env/envs/multi_env_snapshot_local_test.py new file mode 100644 index 000000000..f5517b216 --- /dev/null +++ b/tst/integration/env/envs/multi_env_snapshot_local_test.py @@ -0,0 +1,63 @@ +"""snapshot_env on the local backend, end to end, through a parts grant. + +A MultiEnv's items server takes ``write_object``, so the snapshot hands it an object of its own on the local store; the +server uploads its bundle from its container, the snapshot registers that object as the server's export, and loading the +snapshot back restores the state it held. Requires a docker daemon; spins up a throwaway ``registry:2`` and skips if it +can't start. +""" + +import io +import json +import logging +import uuid +import zipfile + +import pytest +from agentenv_protocol import DataPart +from agentenv_protocol import client as protocol_v1 + +from agent_env.artifact import EnvironmentUniverseArtifact +from agent_env.env import legacy_protocol +from agent_env.env.env import Env +from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.env.envs.multi_env import MultiEnv +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps import snapshot_env as snapshot_mod +from agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep +from tst.integration.env.envs.server_env_local_test import _put_items_image +from tst.integration.store.local_parallel_deploy_test import _docker, _seed_bootstrap_envs, local_stack # noqa: F401 (fixture) + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + + +@pytest.mark.asyncio +async def test_a_snapshot_of_a_server_that_takes_parts_registers_its_upload_and_loads_back(local_stack, caplog): + _seed_bootstrap_envs() + uid = uuid.uuid4().hex[:8] + items = MCPServerEnv.put(id=f"snap-items-{uid}", environment_name="items", + docker_image_artifact=_put_items_image(f"snap-items-{uid}")) + env = MultiEnv.put(id=f"snap-{uid}", mcp_server_envs=[items], name="suite") + + deployed = await env.deploy(sandbox_type="local", ttl_seconds=900) + work_dir = env._sandbox.work_dir + try: + base_url = await legacy_protocol.v1_base_url(deployed, deployed.gateway_url, "items") + await protocol_v1.add_data(base_url, [DataPart(data={"items": ["snap-x", "snap-y"]})]) + + step = SnapshotEnvTaskStep(id="snap", version=None, snapshot_id=f"snap-universe-{uid}") + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + context = await step.execute(TaskStepContext(deployed_envs=[deployed])) + + assert "items exported through parts" in caplog.text + universe = EnvironmentUniverseArtifact.get(context.metadata["env_snapshotted_universes"]["snap"]["id"]) + (bundle,) = universe.get_file_artifacts().values() + assert "/agentenv-snapshots/" in bundle.object_url and bundle.object_url.endswith("/items.zip") + assert (bundle.filename, bundle.content_type) == ("items.zip", "application/zip") + with zipfile.ZipFile(io.BytesIO(bundle.load())) as zf: + assert json.loads(zf.read("items.json")) == {"items": ["snap-x", "snap-y"]} + + await protocol_v1.add_data(base_url, [DataPart(data={"items": ["after-the-snapshot"]})]) + await (await Env.from_instance_id(deployed.instance_id)).load_environment_universe_artifact(universe) + assert (await protocol_v1.get_data(base_url)).model_dump()["parts"][0]["data"] == {"items": ["snap-x", "snap-y"]} + finally: + _docker("compose", "--project-directory", str(work_dir), "down", "-v", "--remove-orphans") diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 9600c5f1c..19f8a69db 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -9,6 +9,7 @@ import re import socket from datetime import datetime +from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -37,6 +38,7 @@ _redact_compose_secrets, ) from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider +from agent_env.providers.sandbox_providers.local_sandbox import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.env_state import LocalPostgresStateProvider @@ -817,6 +819,28 @@ def test_compose_maps_extra_hosts_into_the_gateway(extra_hosts): assert mapped == ({GATEWAY_SERVICE_NAME: list(extra_hosts)} if extra_hosts else {}) +@pytest.mark.parametrize("trust_dir", [None, Path("/state/tls/trust/ab12")], ids=["no-local-grants", "local-grants"]) +def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(trust_dir): + """Services on a local compose stack reach the local grant server over TLS, so they mount its CA's trust files.""" + from agent_env.env.envs.service_db import ServiceDBConfig + + compose = EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack"), + MCPServerConfig(image="mcp-items", environment_name="items")], + gateway_image="agent-gateway", + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + trust_dir=trust_dir, + ) + + services = yaml.safe_load(compose)["services"] + for name in ("slack", "items"): + env = dict(entry.split("=", 1) for entry in services[name]["environment"]) + trusted = {key: env.get(key) for key in LOCAL_TRUST_ENV} + assert trusted == (LOCAL_TRUST_ENV if trust_dir else dict.fromkeys(LOCAL_TRUST_ENV)) + assert services[name].get("volumes") == ([f"{trust_dir}:{LOCAL_TRUST_DIR}:ro"] if trust_dir else None) + + # --- sidecar rendering -------------------------------------------------------- class _FakeImageArtifact: From 797fa0115eac49b5d799cdc70321177817bc1035 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 15:59:53 -0700 Subject: [PATCH 05/11] test(integration): close the grant server the compose fixture's tests start Co-Authored-By: Claude Opus 5.5 (1M context) --- tst/integration/store/local_parallel_deploy_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tst/integration/store/local_parallel_deploy_test.py b/tst/integration/store/local_parallel_deploy_test.py index e31738d32..4451d0468 100644 --- a/tst/integration/store/local_parallel_deploy_test.py +++ b/tst/integration/store/local_parallel_deploy_test.py @@ -26,6 +26,7 @@ from agent_env.task import Task from agent_env.task.store import get_task_instance_store from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep +from agent_env.store.object_store.local.grant_server import grant_server from tst.task import journal_invariants as journal pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -81,6 +82,9 @@ def local_stack(monkeypatch, tmp_path): finally: for cid in _docker("ps", "-aq", "--filter", "name=agent-env-local-").stdout.split(): _docker("rm", "-f", cid) + # Grants handed out here came from this process's grant server, whose certificate is from this test's state + # root; a later test's containers trust another root's CA, so it must start afresh. + grant_server(None, None).close() reset_artifact_store() reset_config() _docker("rm", "-f", reg_name) From f48eba4b426c166a2e2a66615369ce04ffc19fc0 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 16:00:02 -0700 Subject: [PATCH 06/11] test(integration): sort the grant_server import Co-Authored-By: Claude Opus 5.5 (1M context) --- tst/integration/store/local_parallel_deploy_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tst/integration/store/local_parallel_deploy_test.py b/tst/integration/store/local_parallel_deploy_test.py index 4451d0468..ea01a82e3 100644 --- a/tst/integration/store/local_parallel_deploy_test.py +++ b/tst/integration/store/local_parallel_deploy_test.py @@ -23,10 +23,10 @@ from agent_env.config import configure, reset_config, set_image_store from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.store.image_store import LocalRegistryImageStore +from agent_env.store.object_store.local.grant_server import grant_server from agent_env.task import Task from agent_env.task.store import get_task_instance_store from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep -from agent_env.store.object_store.local.grant_server import grant_server from tst.task import journal_invariants as journal pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] From 0508b2ee47cba5f41728c7769de81b706cc45b51 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 18:23:03 -0700 Subject: [PATCH 07/11] fix(local): trusted compose services also carry the sandbox's host mappings On Linux, host.docker.internal resolves in a container only through extra_hosts, which the compose gave the gateway alone. A service handed the local CA's trust files now gets the mapping as well, so it reaches the grant server it trusts. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../providers/env_providers/env_gateway_provider.py | 4 +++- .../providers/env_providers/env_gateway_provider_test.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index aceb4a10d..e4cd65f6e 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -221,7 +221,7 @@ def create_docker_compose( ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. With ``trust_dir``, the local transfer CA's trust files, each MCP server mounts them where - ``LOCAL_TRUST_ENV`` points, so its TLS clients trust the local grant server. + ``LOCAL_TRUST_ENV`` points and carries ``extra_hosts``, so it reaches and trusts the local grant server. Supports MCP servers, websites, or both. Website containers (backend + frontend) are added to the same docker network. The gateway receives WEBSITE_URLS env var for discovery. @@ -369,6 +369,8 @@ def create_docker_compose( if trust_dir is not None: lines.extend(f" - {key}={value}" for key, value in LOCAL_TRUST_ENV.items()) lines.extend([" volumes:", f' - "{trust_dir}:{LOCAL_TRUST_DIR}:ro"']) + if extra_hosts: + lines.extend([" extra_hosts:", *(f' - "{entry}"' for entry in extra_hosts)]) lines.extend([ " healthcheck:", f' test: ["CMD-SHELL", "python3 -c \\"import socket; s=socket.create_connection((\'localhost\',{gateway_port}),2); s.close()\\""]', diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 19f8a69db..e861fe84f 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -821,7 +821,8 @@ def test_compose_maps_extra_hosts_into_the_gateway(extra_hosts): @pytest.mark.parametrize("trust_dir", [None, Path("/state/tls/trust/ab12")], ids=["no-local-grants", "local-grants"]) def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(trust_dir): - """Services on a local compose stack reach the local grant server over TLS, so they mount its CA's trust files.""" + """Services on a local compose stack reach the local grant server over TLS, so they mount its CA's trust files + and carry the sandbox's host mappings, which on Linux are how host.docker.internal resolves.""" from agent_env.env.envs.service_db import ServiceDBConfig compose = EnvironmentGatewayProvider().create_docker_compose( @@ -830,6 +831,7 @@ def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(t gateway_image="agent-gateway", state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), state_instance=LocalPostgresStateProvider.default_instance(), + extra_hosts=("host.docker.internal:host-gateway",), trust_dir=trust_dir, ) @@ -839,6 +841,7 @@ def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(t trusted = {key: env.get(key) for key in LOCAL_TRUST_ENV} assert trusted == (LOCAL_TRUST_ENV if trust_dir else dict.fromkeys(LOCAL_TRUST_ENV)) assert services[name].get("volumes") == ([f"{trust_dir}:{LOCAL_TRUST_DIR}:ro"] if trust_dir else None) + assert services[name].get("extra_hosts") == (["host.docker.internal:host-gateway"] if trust_dir else None) # --- sidecar rendering -------------------------------------------------------- From b18b88f9ddc67cf619e95d611082f4e531376e2d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 18:41:50 -0700 Subject: [PATCH 08/11] fix(local): a local compose mounts the trust files from a copy in its work dir The local CA's folder lives under the state dir, which can sit outside what Docker shares (an XDG_STATE_HOME elsewhere), so mounting it from there could fail. LocalSandbox.stage_local_trust copies the trust files into the sandbox's work dir, which is kept on a shared path, and the compose mounts that copy by a path relative to itself. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../env_providers/env_gateway_provider.py | 16 ++++++++-------- .../providers/sandbox_providers/local_sandbox.py | 11 +++++++++++ .../env_providers/env_gateway_provider_test.py | 13 ++++++------- .../sandbox_providers/local_sandbox_test.py | 11 +++++++++++ 4 files changed, 36 insertions(+), 15 deletions(-) diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index e4cd65f6e..825a8e490 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -10,7 +10,6 @@ import time from dataclasses import dataclass, field from datetime import datetime, timezone -from pathlib import Path from typing import TYPE_CHECKING, Optional from agent_env.attribution import Attribution @@ -34,7 +33,7 @@ GATEWAY_SERVICE_NAME, ) from agent_env.providers.env_providers.env_provider import _SandboxEnvironmentProvider, _size_kwargs, _tool_names -from agent_env.providers.sandbox_providers.local_sandbox import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV, LocalSandbox, local_grant_trust +from agent_env.providers.sandbox_providers.local_sandbox import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV, LocalSandbox from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, port_bindings from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SandboxProvider @@ -217,11 +216,12 @@ def create_docker_compose( host_ips: tuple[str, ...] = (), extra_hosts: tuple[str, ...] = (), mcp_server_name: str | None = None, - trust_dir: Path | None = None, + trust_files: str | None = None, ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. - With ``trust_dir``, the local transfer CA's trust files, each MCP server mounts them where - ``LOCAL_TRUST_ENV`` points and carries ``extra_hosts``, so it reaches and trusts the local grant server. + With ``trust_files``, where the compose finds the local transfer CA's trust files, each MCP server mounts + them where ``LOCAL_TRUST_ENV`` points and carries ``extra_hosts``, so it reaches and trusts the local grant + server. Supports MCP servers, websites, or both. Website containers (backend + frontend) are added to the same docker network. The gateway receives WEBSITE_URLS env var for discovery. @@ -366,9 +366,9 @@ def create_docker_compose( if server.extra_env_vars: for key, value in server.extra_env_vars.items(): lines.append(f" - {key}={value}") - if trust_dir is not None: + if trust_files is not None: lines.extend(f" - {key}={value}" for key, value in LOCAL_TRUST_ENV.items()) - lines.extend([" volumes:", f' - "{trust_dir}:{LOCAL_TRUST_DIR}:ro"']) + lines.extend([" volumes:", f' - "{trust_files}:{LOCAL_TRUST_DIR}:ro"']) if extra_hosts: lines.extend([" extra_hosts:", *(f' - "{entry}"' for entry in extra_hosts)]) lines.extend([ @@ -747,7 +747,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, - trust_dir=await asyncio.to_thread(local_grant_trust) if isinstance(sandbox, LocalSandbox) else None, + trust_files=await asyncio.to_thread(sandbox.stage_local_trust) if isinstance(sandbox, LocalSandbox) else None, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index 1a585075f..e9870a325 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -75,6 +75,8 @@ def _runs_in_container(cmd: list[str]) -> bool: # Where a container finds the local transfer CA's trust files, and the variables that point TLS clients at them: # SSL_CERT_FILE replaces a client's roots, so it gets the public roots plus the CA; NODE_EXTRA_CA_CERTS adds. LOCAL_TRUST_DIR = "/etc/agentenv" +# Where in a sandbox's work dir a compose there finds a copy of those trust files to mount. +_STAGED_TRUST_DIR = "agentenv-trust" LOCAL_TRUST_ENV = { "SSL_CERT_FILE": f"{LOCAL_TRUST_DIR}/ca-bundle.pem", "REQUESTS_CA_BUNDLE": f"{LOCAL_TRUST_DIR}/ca-bundle.pem", @@ -204,6 +206,15 @@ def owns_container(self) -> bool: def work_dir(self) -> Path: return self._work_dir + def stage_local_trust(self) -> str | None: + """Copy the local transfer CA's trust files into the work dir, which Docker shares, for a compose there to + mount; their path relative to the work dir, or None where the store hands out no local grants.""" + trust_dir = local_grant_trust() + if trust_dir is None: + return None + shutil.copytree(trust_dir, self._work_dir / _STAGED_TRUST_DIR, dirs_exist_ok=True) + return f"./{_STAGED_TRUST_DIR}" + def _rewrite_app_arg(self, arg: str) -> str: if arg == "/app" or arg.startswith(("/app/", "/app:")): return str(self._work_dir) + arg[4:] diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index e861fe84f..15a30b8f2 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -9,7 +9,6 @@ import re import socket from datetime import datetime -from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -819,8 +818,8 @@ def test_compose_maps_extra_hosts_into_the_gateway(extra_hosts): assert mapped == ({GATEWAY_SERVICE_NAME: list(extra_hosts)} if extra_hosts else {}) -@pytest.mark.parametrize("trust_dir", [None, Path("/state/tls/trust/ab12")], ids=["no-local-grants", "local-grants"]) -def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(trust_dir): +@pytest.mark.parametrize("trust_files", [None, "./agentenv-trust"], ids=["no-local-grants", "local-grants"]) +def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(trust_files): """Services on a local compose stack reach the local grant server over TLS, so they mount its CA's trust files and carry the sandbox's host mappings, which on Linux are how host.docker.internal resolves.""" from agent_env.env.envs.service_db import ServiceDBConfig @@ -832,16 +831,16 @@ def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(t state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), state_instance=LocalPostgresStateProvider.default_instance(), extra_hosts=("host.docker.internal:host-gateway",), - trust_dir=trust_dir, + trust_files=trust_files, ) services = yaml.safe_load(compose)["services"] for name in ("slack", "items"): env = dict(entry.split("=", 1) for entry in services[name]["environment"]) trusted = {key: env.get(key) for key in LOCAL_TRUST_ENV} - assert trusted == (LOCAL_TRUST_ENV if trust_dir else dict.fromkeys(LOCAL_TRUST_ENV)) - assert services[name].get("volumes") == ([f"{trust_dir}:{LOCAL_TRUST_DIR}:ro"] if trust_dir else None) - assert services[name].get("extra_hosts") == (["host.docker.internal:host-gateway"] if trust_dir else None) + assert trusted == (LOCAL_TRUST_ENV if trust_files else dict.fromkeys(LOCAL_TRUST_ENV)) + assert services[name].get("volumes") == ([f"{trust_files}:{LOCAL_TRUST_DIR}:ro"] if trust_files else None) + assert services[name].get("extra_hosts") == (["host.docker.internal:host-gateway"] if trust_files else None) # --- sidecar rendering -------------------------------------------------------- diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_test.py index f98ff7c8d..372a78104 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_test.py @@ -1165,3 +1165,14 @@ def fail(*args): _HOST_IPS = ls._host_ips + + +def test_a_compose_on_a_local_sandbox_mounts_a_copy_of_the_trust_files_its_work_dir_holds(tmp_path, monkeypatch): + """The CA's own folder may sit outside what Docker shares; the work dir is kept on a shared path.""" + sandbox = LocalSandbox(work_dir=tmp_path / "work") + assert sandbox.stage_local_trust() == "./agentenv-trust" + staged = tmp_path / "work" / "agentenv-trust" + assert {p.name: p.read_bytes() for p in staged.iterdir()} == {p.name: p.read_bytes() for p in local_ca().trust_dir.iterdir()} + + monkeypatch.setattr(ls, "local_grant_trust", lambda: None) + assert LocalSandbox(work_dir=tmp_path / "other").stage_local_trust() is None From b98bdc05fed9d21ced7a909b39ef13b912bc5a6c Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 23:16:16 -0700 Subject: [PATCH 09/11] fix(local): stage the compose's trust files on any sandbox on this machine A local VM's compose services are handed local grants as local Docker's are, so the trust files are written through the sandbox's own host beside the compose file, for every sandbox that runs on this machine, instead of copied into local Docker's work dir alone. A server's own variables win over the trust ones, as they do for a container, and a snapshot reload's compose keeps the files mounted. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/envs/multi_env.py | 1 + .../env_providers/env_gateway_provider.py | 27 +++++-- .../env/envs/multi_env_snapshot_local_test.py | 2 +- .../test_multi_env_snapshot_reload_name.py | 20 ++++++ .../env_gateway_provider_test.py | 71 ++++++++++++++++++- 5 files changed, 112 insertions(+), 9 deletions(-) diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 0e2e3758e..134d04964 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -632,6 +632,7 @@ async def _load_from_snapshot(self, snapshot) -> None: host_ips=self._sandbox.host_ips, extra_hosts=self._sandbox.extra_hosts, mcp_server_name=self._mcp_server_name or self.name, + trust_files=await gw._stage_trust_files(self._sandbox), ) compose_content = compose_content.replace( " volumes:\n - ./init-schemas.sql:/docker-entrypoint-initdb.d/init-schemas.sql\n", "", diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 71a50bf86..a16f62497 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -80,6 +80,8 @@ _LOG_LINE_CHARS = 300 # Gathering it may stall with the connection; past this, the deploy raises the error it had. _DOWN_SERVICES_TIMEOUT_S = 60.0 +# Where beside the compose file its services' copy of the local transfer CA's trust files is kept. +_TRUST_FILES_DIR = "trust" # MCP server health-check budget. Some servers wait on an upstream before they bind; a server # that is up passes its first probe, so the long start period only delays failing a broken one. @@ -279,8 +281,8 @@ def create_docker_compose( ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. With ``trust_files``, where the compose finds the local transfer CA's trust files, each MCP server mounts - them where ``LOCAL_TRUST_ENV`` points and carries ``extra_hosts``, so it reaches and trusts the local grant - server. + them where ``LOCAL_TRUST_ENV`` points, unless its own variables say otherwise, and carries ``extra_hosts``, so + it reaches and trusts the local grant server. Supports MCP servers, websites, or both. Website containers (backend + frontend) are added to the same docker network. The gateway receives WEBSITE_URLS env var for discovery. @@ -427,11 +429,9 @@ def create_docker_compose( f" - ENVIRONMENT_NAME={server.environment_name}", f" - DATABASE_URL={db_url}", ]) - if server.extra_env_vars: - for key, value in server.extra_env_vars.items(): - lines.append(f" - {key}={value}") + server_env = {**(LOCAL_TRUST_ENV if trust_files is not None else {}), **(server.extra_env_vars or {})} + lines.extend(f" - {key}={value}" for key, value in server_env.items()) if trust_files is not None: - lines.extend(f" - {key}={value}" for key, value in LOCAL_TRUST_ENV.items()) lines.extend([" volumes:", f' - "{trust_files}:{LOCAL_TRUST_DIR}:ro"']) if extra_hosts: lines.extend([" extra_hosts:", *(f' - "{entry}"' for entry in extra_hosts)]) @@ -813,6 +813,7 @@ async def _stand_up(_spec): self._all_environments(mcp_servers, website_configs), instance=state_instance ) + trust_files = await self._stage_trust_files(sandbox) compose_content = self.create_docker_compose( mcp_servers=mcp_servers, gateway_image=gateway_env.docker_image_artifact.image_name, @@ -826,7 +827,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, - trust_files=await asyncio.to_thread(sandbox.stage_local_trust) if isinstance(sandbox, LocalSandbox) else None, + trust_files=trust_files, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") @@ -1136,6 +1137,18 @@ async def read_trajectory(self, sandbox: VmSandbox) -> list[dict]: events.append(json.loads(line)) return events + @staticmethod + async def _stage_trust_files(sandbox: VmSandbox) -> str | None: + """Write the local transfer CA's trust files beside the compose file and return where it finds them, on a sandbox + on this machine when the store hands out local grants, which its services are given; None otherwise.""" + if not sandbox.ON_THIS_MACHINE or (trust_dir := await asyncio.to_thread(local_grant_trust)) is None: + return None + directory = f"{GATEWAY_APP_DIR}/{_TRUST_FILES_DIR}" + await sandbox.exec_script(f"mkdir -p {directory}", max_retries=2) + for path in sorted(trust_dir.iterdir()): + await sandbox.write_host_file(await asyncio.to_thread(path.read_bytes), f"{directory}/{path.name}") + return f"./{_TRUST_FILES_DIR}" + @staticmethod async def _write_file(sandbox: VmSandbox, path: str, content: str) -> None: """Write ``content`` and a newline to ``path`` on the sandbox host, the content one quoted argument so no byte diff --git a/tst/integration/env/envs/multi_env_snapshot_local_test.py b/tst/integration/env/envs/multi_env_snapshot_local_test.py index f5517b216..4b2ef71f6 100644 --- a/tst/integration/env/envs/multi_env_snapshot_local_test.py +++ b/tst/integration/env/envs/multi_env_snapshot_local_test.py @@ -25,7 +25,7 @@ from agent_env.task_step.task_steps import snapshot_env as snapshot_mod from agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep from tst.integration.env.envs.server_env_local_test import _put_items_image -from tst.integration.store.local_parallel_deploy_test import _docker, _seed_bootstrap_envs, local_stack # noqa: F401 (fixture) +from tst.integration.store.local_parallel_deploy_test import _docker, _seed_bootstrap_envs, local_stack as local_stack pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] diff --git a/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py b/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py index 3ee8f8da2..0b11024a9 100644 --- a/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py +++ b/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py @@ -41,6 +41,26 @@ async def test_snapshot_reload_keeps_the_env_name(declared, deployed_as, expecte assert compose.call_args.kwargs["host_port"] == env._sandbox.host_port +@pytest.mark.asyncio +async def test_snapshot_reload_keeps_its_services_trusting_the_local_transfer_ca(): + env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name="crm") + env._sandbox = MagicMock(mode="vm", load_docker_images=AsyncMock(), host_ips=("127.0.0.1",)) + db_image = MagicMock(spec=DockerImageArtifact, image_name="snap:1") + snapshot = MagicMock(db_image_artifact_id="snap", db_image_artifact_version=1) + + with patch("agent_env.artifact.Artifact.get", return_value=db_image), \ + patch("agent_env.env.env.Env.get", return_value=MagicMock()), \ + patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id", default_service_db_env_id="db-id")), \ + patch("agent_env.providers.env_state.LocalPostgresStateProvider"), \ + patch.object(EnvironmentGatewayProvider, "_stage_trust_files", AsyncMock(return_value="./trust")) as stage, \ + patch.object(EnvironmentGatewayProvider, "create_docker_compose", side_effect=_Stop) as compose: + with pytest.raises(_Stop): + await env._load_from_snapshot(snapshot) + + stage.assert_awaited_once_with(env._sandbox) + assert compose.call_args.kwargs["trust_files"] == "./trust" + + @pytest.mark.asyncio async def test_snapshot_reload_waits_for_the_gateway_where_it_is_published(): env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name="crm") diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 06abcab84..3fa625a25 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -41,7 +41,7 @@ ) from agent_env.providers.env_providers.env_provider import DeployRefused from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider -from agent_env.providers.sandbox_providers.local_docker.provider import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV +from agent_env.providers.sandbox_providers.local_docker.provider import LOCAL_TRUST_DIR, LOCAL_TRUST_ENV, LocalSandbox from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.sandbox_providers.sandbox_provider import Accepts @@ -1124,6 +1124,75 @@ def test_compose_hands_each_mcp_server_the_local_transfer_ca_when_there_is_one(t assert services[name].get("extra_hosts") == (["host.docker.internal:host-gateway"] if trust_files else None) +def test_a_servers_own_variables_win_over_the_local_transfer_ca(): + from agent_env.env.envs.service_db import ServiceDBConfig + + compose = EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-items", environment_name="items", + extra_env_vars={"SSL_CERT_FILE": "/srv/own-roots.pem"})], + gateway_image="agent-gateway", + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + trust_files="./trust", + ) + + keys = [entry.split("=", 1)[0] for entry in yaml.safe_load(compose)["services"]["items"]["environment"]] + env = dict(entry.split("=", 1) for entry in yaml.safe_load(compose)["services"]["items"]["environment"]) + assert len(keys) == len(set(keys)) + assert env["SSL_CERT_FILE"] == "/srv/own-roots.pem" + assert env["NODE_EXTRA_CA_CERTS"] == LOCAL_TRUST_ENV["NODE_EXTRA_CA_CERTS"] + + +class _HostSandbox: + """Records what a deploy writes on its host.""" + + def __init__(self, on_this_machine: bool) -> None: + self.ON_THIS_MACHINE = on_this_machine + self.scripts: list[str] = [] + self.files: dict[str, bytes] = {} + + async def exec_script(self, script: str, **_) -> str: + self.scripts.append(script) + return "" + + async def write_host_file(self, data: bytes, vm_path: str) -> None: + self.files[vm_path] = data + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("on_this_machine", "local_grants", "staged"), [ + (True, True, "./trust"), + (True, False, None), + (False, True, None), +], ids=["on-this-machine", "no-local-grants", "elsewhere"]) +async def test_a_deploy_on_this_machine_writes_the_trust_files_beside_its_compose( + tmp_path, monkeypatch, on_this_machine, local_grants, staged +): + """Any sandbox on this machine, a local VM's as well as local Docker's, gets them through its own host, which its + compose mounts from; a sandbox elsewhere is handed no local grants, so it gets none.""" + (tmp_path / "ca.pem").write_bytes(b"ca") + (tmp_path / "ca-bundle.pem").write_bytes(b"roots+ca") + monkeypatch.setattr(env_gateway_provider, "local_grant_trust", lambda: tmp_path if local_grants else None) + sandbox = _HostSandbox(on_this_machine) + + assert await EnvironmentGatewayProvider._stage_trust_files(sandbox) == staged + + expected = {"/app/trust/ca-bundle.pem": b"roots+ca", "/app/trust/ca.pem": b"ca"} if staged else {} + assert sandbox.files == expected + assert sandbox.scripts == (["mkdir -p /app/trust"] if staged else []) + + +@pytest.mark.asyncio +async def test_a_local_sandbox_keeps_the_trust_files_in_its_work_dir(tmp_path, monkeypatch): + """The CA's own folder may sit outside what Docker shares; the work dir, where the compose runs, is shared.""" + (tmp_path / "ca").mkdir() + (tmp_path / "ca" / "ca.pem").write_bytes(b"ca") + monkeypatch.setattr(env_gateway_provider, "local_grant_trust", lambda: tmp_path / "ca") + + assert await EnvironmentGatewayProvider._stage_trust_files(LocalSandbox(work_dir=tmp_path / "work")) == "./trust" + assert (tmp_path / "work" / "trust" / "ca.pem").read_bytes() == b"ca" + + # --- sidecar rendering -------------------------------------------------------- class _FakeImageArtifact: From b74b18944d8ab6068fa6f6ec86f0215fa7a1098c Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 23:28:00 -0700 Subject: [PATCH 10/11] fix(sandbox): decode a chunked host write from stdin, which macOS's base64 also reads A host write past one chunk decoded `base64 -d FILE`, which macOS's base64 refuses, so a local sandbox on a Mac failed any large write_host_file, as staging a compose's trust bundle does on every local MultiEnv deploy. Reading the file from stdin decodes the same on Linux and macOS. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/providers/sandbox_providers/sandbox.py | 3 ++- .../providers/env_providers/env_gateway_provider_test.py | 9 ++++++--- tst/unit/providers/sandbox_providers/vm_sandbox_test.py | 4 ++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 86c740656..711032afd 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -597,7 +597,8 @@ async def _write_bytes_to_vm_path(self, data: bytes, vm_path: str) -> None: await self.exec_script(f": > {shlex.quote(vm_b64)}") for i in range(0, len(encoded), self._WFT_CHUNK_BYTES): await self.exec_script(f"printf '%s' {shlex.quote(encoded[i:i + self._WFT_CHUNK_BYTES])} >> {shlex.quote(vm_b64)}") - await self.exec_script(f"base64 -d {shlex.quote(vm_b64)} > {shlex.quote(vm_path)} && rm -f {shlex.quote(vm_b64)}") + # From stdin: macOS's base64, a local sandbox's, takes no file argument. + await self.exec_script(f"base64 -d < {shlex.quote(vm_b64)} > {shlex.quote(vm_path)} && rm -f {shlex.quote(vm_b64)}") async def write_host_file(self, data: bytes, vm_path: str) -> None: """Write bytes to vm_path on the VM host itself, not into the agent container.""" diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 3fa625a25..0f864fcb0 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -1184,13 +1184,16 @@ async def test_a_deploy_on_this_machine_writes_the_trust_files_beside_its_compos @pytest.mark.asyncio async def test_a_local_sandbox_keeps_the_trust_files_in_its_work_dir(tmp_path, monkeypatch): - """The CA's own folder may sit outside what Docker shares; the work dir, where the compose runs, is shared.""" + """The CA's own folder may sit outside what Docker shares; the work dir, where the compose runs, is shared. The bundle + of public roots and the CA is a few hundred KB, so it is written in chunks.""" (tmp_path / "ca").mkdir() - (tmp_path / "ca" / "ca.pem").write_bytes(b"ca") + trust = {"ca.pem": b"ca", "ca-bundle.pem": b"r" * (300 * 1024)} + for name, content in trust.items(): + (tmp_path / "ca" / name).write_bytes(content) monkeypatch.setattr(env_gateway_provider, "local_grant_trust", lambda: tmp_path / "ca") assert await EnvironmentGatewayProvider._stage_trust_files(LocalSandbox(work_dir=tmp_path / "work")) == "./trust" - assert (tmp_path / "work" / "trust" / "ca.pem").read_bytes() == b"ca" + assert {p.name: p.read_bytes() for p in (tmp_path / "work" / "trust").iterdir()} == trust # --- sidecar rendering -------------------------------------------------------- diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py index 2e77e3540..632b79953 100644 --- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py @@ -469,7 +469,7 @@ async def test_write_file_from_text_large_chunks_reconstruct_exactly(): # The reconstructed base64 must equal the original, byte-for-byte. assert _b64_from_chunk_scripts(sandbox.scripts, b64_path) == expected_b64 # And it's decoded into the real path before the docker cp. - assert any(f"base64 -d {b64_path} > " in s for s in sandbox.scripts) + assert any(f"base64 -d < {b64_path} > " in s for s in sandbox.scripts) class _CleanupFailingSandbox(VmSandbox): @@ -715,7 +715,7 @@ async def test_write_host_file_writes_on_the_host_in_bounded_chunks(): assert not any("docker" in s for s in sandbox.scripts) b64_path = "/tmp/agentenv_run_code/input.json.b64" assert _b64_from_chunk_scripts(sandbox.scripts, b64_path) == base64.b64encode(data).decode() - assert sandbox.scripts[-1] == f"base64 -d {b64_path} > /tmp/agentenv_run_code/input.json && rm -f {b64_path}" + assert sandbox.scripts[-1] == f"base64 -d < {b64_path} > /tmp/agentenv_run_code/input.json && rm -f {b64_path}" class _ArgsRecorder(VmSandbox): From a7411d1d55c5fa6febaa0412ba9d0172e4297d79 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 23:54:24 -0700 Subject: [PATCH 11/11] fix(snapshot_env): log a write a store made through one PUT, and make staged trust files readable A store may answer a service that takes parts with one PUT; the export's log line no longer reads a part size that grant lacks, which failed the whole snapshot after the object was made. The trust files staged for a compose are made 755/644 whatever the deploy's umask, so a service running as any user reads them, and the deploy's mount of them is pinned in a unit test. Co-Authored-By: Claude Opus 5.5 --- .../env_providers/env_gateway_provider.py | 8 +-- .../task_step/task_steps/snapshot_env.py | 6 ++- .../env_gateway_provider_test.py | 53 +++++++++++++++++-- .../snapshot_env_object_write_test.py | 34 ++++++++++++ 4 files changed, 92 insertions(+), 9 deletions(-) diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index a16f62497..dac3bc46a 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -1140,13 +1140,15 @@ async def read_trajectory(self, sandbox: VmSandbox) -> list[dict]: @staticmethod async def _stage_trust_files(sandbox: VmSandbox) -> str | None: """Write the local transfer CA's trust files beside the compose file and return where it finds them, on a sandbox - on this machine when the store hands out local grants, which its services are given; None otherwise.""" + on this machine when the store hands out local grants, which its services are given; None otherwise. They are + made readable by any user a service runs as, whatever the umask they were written under.""" if not sandbox.ON_THIS_MACHINE or (trust_dir := await asyncio.to_thread(local_grant_trust)) is None: return None - directory = f"{GATEWAY_APP_DIR}/{_TRUST_FILES_DIR}" + directory = shlex.quote(f"{GATEWAY_APP_DIR}/{_TRUST_FILES_DIR}") await sandbox.exec_script(f"mkdir -p {directory}", max_retries=2) for path in sorted(trust_dir.iterdir()): - await sandbox.write_host_file(await asyncio.to_thread(path.read_bytes), f"{directory}/{path.name}") + await sandbox.write_host_file(await asyncio.to_thread(path.read_bytes), f"{GATEWAY_APP_DIR}/{_TRUST_FILES_DIR}/{path.name}") + await sandbox.exec_script(f"chmod 755 {directory} && chmod 644 {directory}/*", max_retries=2) return f"./{_TRUST_FILES_DIR}" @staticmethod diff --git a/src/agent_env/task_step/task_steps/snapshot_env.py b/src/agent_env/task_step/task_steps/snapshot_env.py index e05311266..dc696a29a 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -19,7 +19,7 @@ from agentenv_protocol import DATA_OBJECTS_EXTENSION_URI, FilePart, uploaded_object, uploaded_object_path from agentenv_protocol.client import find_extension, find_extension_method, write_kinds -from agentenv_protocol.transfers import WriteNamespaceGrant, part_ranges +from agentenv_protocol.transfers import HttpPartsPutGrant, WriteNamespaceGrant, part_ranges from agent_env.a2a_agent.object_transfer import ObjectLimits, namespace_grant from agent_env.env.gateway.constants import CHANGES_PAGE_MAX, EXT_CHANGES_URI, EXT_TRAJECTORY_URI @@ -553,7 +553,9 @@ def exported(tier: str, result: str, detail: str = "") -> str: part = resp.parts[0] if resp.parts else None if write is not None and (uploaded := uploaded_object(part)) is not None: await _finish_in_thread(write.complete, uploaded) - parts = len(part_ranges(uploaded.size_bytes, write.grant.write.part_bytes)) + if not isinstance(grant := write.grant.write, HttpPartsPutGrant): + return exported("object", write.object_url, f" ({uploaded.size_bytes} bytes)") + parts = len(part_ranges(uploaded.size_bytes, grant.part_bytes)) return exported("parts", write.object_url, f" ({uploaded.size_bytes} bytes in {parts} parts)") finally: if write is not None: diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 0f864fcb0..466a0a809 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -7,9 +7,11 @@ import dataclasses import json import logging +import os import re import shlex import socket +import stat import subprocess from datetime import datetime from types import SimpleNamespace @@ -1179,21 +1181,29 @@ async def test_a_deploy_on_this_machine_writes_the_trust_files_beside_its_compos expected = {"/app/trust/ca-bundle.pem": b"roots+ca", "/app/trust/ca.pem": b"ca"} if staged else {} assert sandbox.files == expected - assert sandbox.scripts == (["mkdir -p /app/trust"] if staged else []) + assert sandbox.scripts == (["mkdir -p /app/trust", "chmod 755 /app/trust && chmod 644 /app/trust/*"] if staged else []) @pytest.mark.asyncio async def test_a_local_sandbox_keeps_the_trust_files_in_its_work_dir(tmp_path, monkeypatch): """The CA's own folder may sit outside what Docker shares; the work dir, where the compose runs, is shared. The bundle - of public roots and the CA is a few hundred KB, so it is written in chunks.""" + of public roots and the CA is a few hundred KB, so it is written in chunks, and a service running as any user reads + the files, whatever the deploy's umask.""" (tmp_path / "ca").mkdir() trust = {"ca.pem": b"ca", "ca-bundle.pem": b"r" * (300 * 1024)} for name, content in trust.items(): (tmp_path / "ca" / name).write_bytes(content) monkeypatch.setattr(env_gateway_provider, "local_grant_trust", lambda: tmp_path / "ca") - assert await EnvironmentGatewayProvider._stage_trust_files(LocalSandbox(work_dir=tmp_path / "work")) == "./trust" - assert {p.name: p.read_bytes() for p in (tmp_path / "work" / "trust").iterdir()} == trust + umask = os.umask(0o077) + try: + assert await EnvironmentGatewayProvider._stage_trust_files(LocalSandbox(work_dir=tmp_path / "work")) == "./trust" + finally: + os.umask(umask) + staged = tmp_path / "work" / "trust" + assert {p.name: p.read_bytes() for p in staged.iterdir()} == trust + assert {p.name: stat.S_IMODE(p.stat().st_mode) for p in [staged, *staged.iterdir()]} == { + "trust": 0o755, "ca.pem": 0o644, "ca-bundle.pem": 0o644} # --- sidecar rendering -------------------------------------------------------- @@ -1453,6 +1463,41 @@ def _fake_env_get(env_id, *a, **k): assert browser.card_port == 18766 # its card is the recorder's +@pytest.mark.asyncio +async def test_deploy_via_vm_mounts_the_trust_files_it_staged_on_its_sandbox(): + from agent_env.env.gateway import GatewayMode + from agent_env.providers.env_state import DatabaseStateProvider + + gp = EnvironmentGatewayProvider() + gp._state_provider = MagicMock(spec=DatabaseStateProvider, prepare=AsyncMock()) + gp._state_instance = MagicMock() + sandbox = AsyncMock() + sandbox_provider = MagicMock(create_vm=AsyncMock(return_value=sandbox)) + + class _Stop(Exception): + """Abort once the compose is rendered.""" + + with patch("agent_env.env.env.Env.get", side_effect=lambda env_id, *a, **k: MagicMock(docker_image_artifact=_loadable(env_id))), \ + patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id")), \ + patch.object(EnvironmentGatewayProvider, "_stage_trust_files", AsyncMock(return_value="./trust")) as stage, \ + patch.object(gp, "create_docker_compose", side_effect=_Stop) as compose: + with pytest.raises(_Stop): + await gp._deploy_via_vm( + sandbox_provider=sandbox_provider, + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], + mcp_server_images=[_loadable("mcp-slack")], + gateway_port=18765, + website_configs=None, + website_images=None, + gateway_mode=GatewayMode.PERFORMANCE, + ttl_seconds=60, + disk_size_gb=10, + ) + + stage.assert_awaited_once_with(sandbox) + assert compose.call_args.kwargs["trust_files"] == "./trust" + + @pytest.mark.asyncio async def test_deploy_via_vm_side_loads_sidecar_image(): """A sidecar's image_artifact is side-loaded onto the VM via load_docker_images alongside diff --git a/tst/unit/task_step/snapshot_env_object_write_test.py b/tst/unit/task_step/snapshot_env_object_write_test.py index adbf93b0d..7b0ca9d65 100644 --- a/tst/unit/task_step/snapshot_env_object_write_test.py +++ b/tst/unit/task_step/snapshot_env_object_write_test.py @@ -50,6 +50,13 @@ def begin_write(self, object_url, **kwargs): return write +class _PutOnlyStore(_WatchedStore): + """Answers a service that takes parts with one PUT, as a store may.""" + + def begin_write(self, object_url, *, kinds, **kwargs): + return super().begin_write(object_url, kinds=kinds - {"http-put-parts"}, **kwargs) + + class _GatedStore(_WatchedStore): """Holds the export at ``at`` (beginning the write, pushing creds, completing it) until released.""" @@ -150,6 +157,33 @@ async def test_a_service_taking_parts_uploads_to_an_object_of_its_own_and_is_reg assert f"gdrive exported through parts ({len(_BUNDLE)} bytes in 1 parts)" in caplog.text +@pytest.mark.asyncio +async def test_a_service_taking_parts_that_is_handed_one_put_is_registered_there_too( + local_stores, tmp_path, pushes, monkeypatch, caplog +): + store = _PutOnlyStore(str(tmp_path / "objects")) + set_object_store(store) + record = await _record({"mcp-gdrive": [_PARTS]}) + + def answer(params): + grant = params["write_object"]["write"] + assert grant["kind"] == "http-put" + store.put(unquote(urlsplit(grant["url"]).path.lstrip("/")), _BUNDLE, content_type="application/zip") + part = uploaded_object_part(Uploaded(size_bytes=len(_BUNDLE)), name="gdrive.zip", mime_type="application/zip") + return {"parts": [part.model_dump(mode="json", exclude_none=True)]} + + _serve(monkeypatch, {"gdrive": answer}) + with caplog.at_level(logging.INFO, logger=snapshot_mod.logger.name): + result = await SnapshotEnvTaskStep.snapshot_env_state( + env=_FakeMultiEnv(["gdrive"]), gateway_url="https://gw-1", snapshot_id="snap-put", deployed=record, + ) + + assert store.ended == ["completed"] + (bundle,) = EnvironmentUniverseArtifact.get(result.environment_universe_artifact_id).get_file_artifacts().values() + assert bundle.load() == _BUNDLE + assert f"gdrive exported through object ({len(_BUNDLE)} bytes)" in caplog.text + + @pytest.mark.asyncio async def test_each_object_write_gets_an_object_of_its_own(store, pushes, monkeypatch): record = await _record({"mcp-gdrive": [_PARTS]})