diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index e716c8a2..6997f342 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -20,6 +20,7 @@ invoke_transfer, skill_add_call, ) +from agent_env.a2a_agent.staging import transfer_store from agent_env.providers.sandbox_providers.local_sandbox import ( LOCAL_TRUST_ENV, LocalSandbox, @@ -197,10 +198,13 @@ async def register_skill( """Send one skill, as SKILL.md text inline or the objects under ``object_url`` as a bundle of read grants; returns the agent's answer.""" add_method, add_path = A2AAgent.operation(A2AAgent._skill_extension(deployed), "add") + store = transfer_store( + get_config().get_object_store(), deployed.a2a_url, deployed.agent_card, sandbox_type=deployed.sandbox_type + ) call = await asyncio.to_thread( skill_add_call, add_method, - get_config().get_object_store(), + store, name=name, description=description, skill_md=skill_md, @@ -213,6 +217,7 @@ async def register_skill( verb="POST", operation=f"skill add ({name})", timeout=TRANSFER_TIMEOUT_SECONDS, + store=store, ) @staticmethod diff --git a/src/agent_env/a2a_agent/object_transfer.py b/src/agent_env/a2a_agent/object_transfer.py index 4736eff8..41b4b0f6 100644 --- a/src/agent_env/a2a_agent/object_transfer.py +++ b/src/agent_env/a2a_agent/object_transfer.py @@ -36,6 +36,7 @@ from pydantic import BaseModel, ValidationError from agent_env.a2a_agent.protocol import raise_for_extension_status +from agent_env.a2a_agent.staging import StagedObjectStore from agent_env.store.base import GrantUnavailableError from agent_env.store.object_store import DEFAULT_CONTENT_TYPE, ObjectStore from agent_env.store.object_store.local.grant_server import unreachable_hint @@ -49,6 +50,11 @@ _MEDIA_TYPE = re.compile(r"[!#$&^_.+\-|~0-9a-z]+/[!#$&^_.+\-|~0-9a-z]+") _INCREMENT_NAME = re.compile(r"^(?P[0-9]{6})(?:\.[A-Za-z0-9][A-Za-z0-9._-]*)?$") _TRANSFER_UNAVAILABLE = "transfer_unavailable" # the SDK's code for a store it could not reach +_STAGING_UNREACHABLE = ( + "The object store's grants do not reach this agent, so its objects were staged on the agent's own " + "server, and the agent could not reach them there through its own URL. Its sandbox provider may not let " + "a sandbox call its own public URL; an object store whose grants reach the agent avoids staging." +) # The time budget of one transfer, outermost first: a grant (the issuing store's # grant_lifetime_seconds, 12 hours unless configured) outlives agent-env's wait for the agent's @@ -163,21 +169,34 @@ async def invoke_transfer( operation: str, timeout: float, response_model: type[_Response] | None = None, + store: ObjectStore | None = None, ) -> Any: """Send ``call`` and return the agent's answer, validated against ``response_model`` when the call carried grants. An error never quotes the answer to such a call, which may echo a - grant.""" - async with httpx.AsyncClient() as client: - send = client.post if verb == "POST" else client.put - resp = await send(url, json=call.payload, timeout=timeout) + grant. ``store`` is the store the call's grants came from: when it stages them on the agent, + the objects the agent reads are pushed first and the ones it writes are in the store on return.""" + staged = store if isinstance(store, StagedObjectStore) and call.mode == "objects" else None try: - raise_for_extension_status(resp, operation=operation, include_body=call.mode == "legacy") - except httpx.HTTPStatusError as exc: - hint = unreachable_hint(call.payload) if call.mode == "objects" else None - if hint is None or not str(exc).endswith(f": {_TRANSFER_UNAVAILABLE}"): - raise - raise httpx.HTTPStatusError(f"{exc}. {hint}", request=exc.request, response=exc.response) from exc - body = resp.json() + if staged is not None: + await staged.push() + async with httpx.AsyncClient() as client: + send = client.post if verb == "POST" else client.put + resp = await send(url, json=call.payload, timeout=timeout) + try: + raise_for_extension_status(resp, operation=operation, include_body=call.mode == "legacy") + except httpx.HTTPStatusError as exc: + hint = _STAGING_UNREACHABLE if staged is not None else ( + unreachable_hint(call.payload) if call.mode == "objects" else None + ) + if hint is None or not str(exc).endswith(f": {_TRANSFER_UNAVAILABLE}"): + raise + raise httpx.HTTPStatusError(f"{exc}. {hint}", request=exc.request, response=exc.response) from exc + body = resp.json() + if staged is not None: + await staged.pull() + finally: + if staged is not None: + await staged.release() if call.mode == "objects" and response_model is not None: return parse_response(response_model, body, operation=operation) return body @@ -538,9 +557,10 @@ async def fetch_trajectory( *, upload: TrajectoryUpload | None = None, timeout: float | None = None, + store: ObjectStore | None = None, ) -> FetchedTrajectory: """POST one trajectory ``get``. With ``upload`` the agent uploads the trajectory through its - grant and the answer is checked. The default wait covers that upload.""" + grant, issued by ``store``, and the answer is checked. The default wait covers that upload.""" payload: dict[str, Any] = dict(selector) if upload is not None: objects = TrajectoryWriteObjects(trajectory=upload.write) @@ -555,6 +575,7 @@ async def fetch_trajectory( operation="trajectory get", timeout=timeout, response_model=TrajectoryObjectsResponse, + store=store, ) if upload is not None: return FetchedTrajectory(object_url=upload.object_url) diff --git a/src/agent_env/a2a_agent/staging.py b/src/agent_env/a2a_agent/staging.py new file mode 100644 index 00000000..d3215106 --- /dev/null +++ b/src/agent_env/a2a_agent/staging.py @@ -0,0 +1,393 @@ +"""Objects moved through an agent's staging extension, for agents the object store's grants cannot reach. + +A local object store's grants reach only agents on this machine. An agent elsewhere that serves +``urn:agentenv:staging/v1`` (every SDK agent does) gets ordinary HTTPS grants naming paths on its own +staging routes instead: agent-env pushes each object the agent reads there before the call, and pulls each +object it writes into the store after it. A changelog namespace stays staged for the agent's life, so its +increments are drained into the store while the agent works, when a prompt ends, and before the agent is +torn down. Every byte moves over a connection agent-env opens, so nothing has to reach this machine. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +import secrets +import tempfile +import threading +from collections.abc import AsyncIterator, Iterable, Mapping +from contextlib import asynccontextmanager +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any, BinaryIO +from urllib.parse import quote, urlsplit + +import httpx +from agentenv_protocol.a2a_agent import STAGING_V1_URI +from agentenv_protocol.transfers import HttpGetGrant, HttpPostPolicyGrant, HttpPutGrant, redacting_request_urls + +from agent_env.config import get_config +from agent_env.store.base import ObjectAlreadyExistsError +from agent_env.store.object_store import DEFAULT_CONTENT_TYPE, ObjectMetadata, ObjectStore, UploadPolicy + +logger = logging.getLogger(__name__) + +# How often a running prompt's staged changelog increments are moved into the store: what is lost if +# the agent's sandbox dies mid-prompt is at most this much of its work. +DRAIN_INTERVAL_SECONDS = 5.0 +# The longest a last drain waits on an agent before giving up on what it still holds, so a stalled agent can't +# hold up the end of a prompt or a teardown. +LAST_DRAIN_SECONDS = 120.0 +_CHUNK_BYTES = 1024 * 1024 +_ATTEMPTS = 3 +_CONCURRENCY = 4 +_TIMEOUT = httpx.Timeout(connect=30.0, read=60.0, write=60.0, pool=30.0) +_REPLACING = threading.Lock() + + +class StagingError(RuntimeError): + """Moving an object through an agent's staging failed. The message never names a staged path, whose id + is what keeps it private.""" + + +def staging_endpoint(a2a_url: str, card: Mapping[str, Any] | None) -> str | None: + """The absolute URL of the agent's staging routes, when its card advertises them on an HTTPS URL a + grant can name; None otherwise.""" + params = _staging_params(card) + endpoint = params.get("endpoint") + url = f"{a2a_url.rstrip('/')}{endpoint}" if isinstance(endpoint, str) and endpoint.startswith("/") else None + return url.rstrip("/") if url is not None and urlsplit(url).scheme == "https" else None + + +def _staging_params(card: Mapping[str, Any] | None) -> Mapping[str, Any]: + for extension in ((card or {}).get("capabilities") or {}).get("extensions") or (): + if isinstance(extension, Mapping) and extension.get("uri") == STAGING_V1_URI: + return extension.get("params") or {} + return {} + + +def transfer_store( + store: ObjectStore, a2a_url: str, card: Mapping[str, Any] | None, *, sandbox_type: str | None +) -> ObjectStore: + """The store an extension call's grants come from: ``store`` when its grants reach the agent, else a + store that stages them through the agent when the agent serves staging, else ``store`` itself, so the + call falls back to the forms that carry no grants.""" + endpoint = staging_endpoint(a2a_url, card) + if endpoint is None or (store.supports_transfer_grants and store.grants_reach(sandbox_type)): + return store + limit = _staging_params(card).get("max_bytes") + return StagedObjectStore(store, endpoint, max_bytes=limit if isinstance(limit, int) and limit > 0 else None) + + +@dataclass(frozen=True) +class StagedNamespace: + """A changelog namespace staged on an agent: where its increments wait, and where they go.""" + + staging_url: str # holds an id no one else knows: kept with the run, never logged + namespace_url: str + max_object_bytes: int | None = None + + +@dataclass(frozen=True) +class _Staged: + path: str + object_url: str + media_type: str = DEFAULT_CONTENT_TYPE + max_bytes: int | None = None # the most the grant let the agent write + + +class StagedObjectStore: + """``store``, issuing grants that name paths on the agent's staging routes at ``endpoint``. + + One serves one extension call: build the call with it, then pass it to ``invoke_transfer``, which + pushes the reads before the call, pulls the writes after it, and clears the call's staging either + way. A namespace it issues outlives the call; ``namespaces`` names them for the caller to keep. + """ + + supports_transfer_grants = True + + def __init__(self, store: ObjectStore, endpoint: str, *, max_bytes: int | None = None) -> None: + self.store = store + self.endpoint = endpoint + self.max_bytes = max_bytes # what the agent's card says its staging holds + self._call = secrets.token_urlsafe(24) + self._reads: list[_Staged] = [] + self._writes: list[_Staged] = [] + self._sizes: dict[str, int | None] = {} # each object's size when the call described it + self.namespaces: list[StagedNamespace] = [] + + def __getattr__(self, name: str) -> Any: + return getattr(self.store, name) + + def grants_reach(self, sandbox_type: str | None) -> bool: + return True + + def get_object_metadata_at(self, object_url: str) -> ObjectMetadata | None: + metadata = self.store.get_object_metadata_at(object_url) + self._sizes[object_url] = None if metadata is None else metadata.size + return metadata + + def issue_read_grant(self, object_url: str, *, expires_in: int | None = None) -> HttpGetGrant: + staged = self._stage(self._reads, object_url) + return HttpGetGrant(kind="http-get", url=self._url(staged.path), expires_at=self._expiry(expires_in)) + + def issue_write_grant( + self, object_url: str, *, media_type: str, max_bytes: int, expires_in: int | None = None + ) -> HttpPutGrant: + staged = self._stage(self._writes, object_url, media_type, max_bytes) + return HttpPutGrant(kind="http-put", url=self._url(staged.path), expires_at=self._expiry(expires_in)) + + def issue_upload_policy(self, prefix_url: str, *, max_object_bytes: int, expires_in: int) -> UploadPolicy: + url = f"{self.endpoint}/{secrets.token_urlsafe(24)}" + self.namespaces.append(StagedNamespace(url, prefix_url, max_object_bytes)) + return UploadPolicy( + write=HttpPostPolicyGrant(kind="http-post-policy", url=url, fields={}, path_field="key", file_field="file"), + expires_at=self._expiry(expires_in), + ) + + async def push(self) -> None: + """Put each object the call reads into the agent's staging, at the size the call described it with; objects + that together outgrow the limit its card advertises are refused before any is sent.""" + sizes = [ + self._sizes[staged.object_url] if staged.object_url in self._sizes + else await asyncio.to_thread(_size, self.store, staged.object_url) + for staged in self._reads + ] + if self.max_bytes is not None and sum(size or 0 for size in sizes) > self.max_bytes: + raise StagingError("staging an object on the agent failed: its staging is full") + async with _client("staging an object on the agent") as client: + await _each( + list(zip(self._reads, sizes)), + lambda item: _push(client, self.store, item[0].object_url, self._url(item[0].path), item[1]), + ) + + async def pull(self) -> None: + """Move each object the agent wrote into the store; one it did not write is skipped.""" + async with _client("collecting an object the agent staged") as client: + await _each(self._writes, lambda staged: _pull(client, self.store, self._url(staged.path), staged)) + + async def release(self) -> None: + """Clear what the call staged. Best-effort: the agent's sandbox ending clears it too.""" + if not self._reads and not self._writes: + return + try: + async with _client("clearing the call's staged objects") as client: + (await client.delete(self._url("") + "/")).raise_for_status() + except StagingError as exc: + logger.warning("%s", exc) + + def _stage( + self, staged: list[_Staged], object_url: str, media_type: str = DEFAULT_CONTENT_TYPE, max_bytes: int | None = None + ) -> _Staged: + entry = _Staged(str(len(self._reads) + len(self._writes)), object_url, media_type, max_bytes) + staged.append(entry) + return entry + + def _url(self, path: str) -> str: + return f"{self.endpoint}/{self._call}" + (f"/{quote(path)}" if path else "") + + def _expiry(self, expires_in: int | None) -> datetime: + seconds = expires_in if expires_in is not None else self.store.grant_lifetime_seconds + return datetime.now(UTC) + timedelta(seconds=seconds) + + +async def drain(namespace: StagedNamespace, store: ObjectStore) -> int: + """Move the increments staged in ``namespace`` into ``store``, and return how many. An increment is + removed from staging only once it is in the store, and only if the agent has not replaced it since.""" + root = store.get_object_key(namespace.namespace_url).rstrip("/") + "/" + async with _client("draining staged changelog increments") as client: + listed = await client.get(namespace.staging_url + "/") + listed.raise_for_status() + entries = [e for e in listed.json().get("objects") or () if str(e.get("path", "")).startswith(root)] + + async def move(entry: Mapping[str, Any]) -> None: + staged = _Staged(entry["path"], store.object_url(entry["path"]), max_bytes=namespace.max_object_bytes) + await _pull(client, store, f"{namespace.staging_url}/{quote(entry['path'])}", staged) + + await _each(entries, move) + return len(entries) + + +def staged_changelogs(metadata: Mapping[str, Any], *, agent_name: str | None = None) -> list[StagedNamespace]: + """The changelog namespaces a run's metadata records as staged on its agents, ``agent_name``'s alone + when one is given.""" + return [ + StagedNamespace(entry["staging_url"], entry["object_url"], entry.get("staging_max_object_bytes")) + for entry in metadata.get("agent_changelog") or () + if entry.get("staging_url") and agent_name in (None, entry.get("agent_name")) + ] + + +async def drain_all(namespaces: Iterable[StagedNamespace], *, within: float | None = None, quiet: bool = False) -> list[str]: + """Drain each of ``namespaces`` into the store its namespace URL is in, giving up after ``within`` + seconds when given, and return why any could not be. One that cannot be drained is left for the next + drain to try again, and logged unless ``quiet``.""" + failures = [] + try: + async with asyncio.timeout(within): + for namespace in namespaces: + try: + await drain(namespace, get_config().get_object_store_at(namespace.namespace_url)) + except Exception as exc: # an agent that is gone or busy must not fail the run + failures.append(str(exc) if isinstance(exc, StagingError) else type(exc).__name__) + except TimeoutError: + failures.append(f"staged changelog increments were not all drained within {within:g}s") + for failure in failures if not quiet else (): + logger.warning("%s", failure) + return failures + + +@asynccontextmanager +async def draining(namespaces: Iterable[StagedNamespace]) -> AsyncIterator[None]: + """Drain ``namespaces`` every ``DRAIN_INTERVAL_SECONDS`` for the length of the block, and once more + when it ends.""" + namespaces = list(namespaces) + if not namespaces: + yield + return + + async def loop() -> None: + reported: list[str] = [] + while True: + await asyncio.sleep(DRAIN_INTERVAL_SECONDS) + failures = await drain_all(namespaces, quiet=True) + if failures and failures != reported: # when the reason changes, not every few seconds + logger.warning("%s; retrying every %gs", "; ".join(failures), DRAIN_INTERVAL_SECONDS) + elif reported and not failures: + logger.info("staged changelog increments are draining again") + reported = failures + + task = asyncio.create_task(loop()) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await drain_all(namespaces, within=LAST_DRAIN_SECONDS) + + +def _size(store: ObjectStore, object_url: str) -> int | None: + metadata = store.get_object_metadata_at(object_url) + return metadata.size if metadata is not None else None + + +async def _push(client: httpx.AsyncClient, store: ObjectStore, object_url: str, url: str, size: int | None) -> None: + async def attempt() -> None: + source = await asyncio.to_thread(store.open, object_url) + headers = {"content-type": "application/octet-stream"} + if size is not None: + headers["content-length"] = str(size) + response = await client.put(url, content=_chunks(source), headers=headers) + response.raise_for_status() + + await _retrying(attempt) + + +async def _pull(client: httpx.AsyncClient, store: ObjectStore, url: str, staged: _Staged) -> None: + fd, name = tempfile.mkstemp(prefix="agentenv-staged-") + try: + with os.fdopen(fd, "wb") as out: + fetched = await _retrying(lambda: _fetch(client, url, out, staged.max_bytes)) + if fetched is None: + return + etag, digest = fetched + try: + await asyncio.to_thread(store.put_file_at, staged.object_url, name, staged.media_type) + except ObjectAlreadyExistsError: + # Rewritten after an earlier copy went into the store: the latest write wins, as through a grant. + await asyncio.to_thread(_replace, store, staged, Path(name), digest) + removed = await client.delete(url, headers={"if-match": etag}) + if removed.status_code not in (204, 404, 412): # 412: rewritten since, so the next drain takes it + removed.raise_for_status() + finally: + os.unlink(name) + + +async def _fetch(client: httpx.AsyncClient, url: str, out: BinaryIO, max_bytes: int | None) -> tuple[str, str] | None: + """Stream the object at ``url`` into ``out``, no more than ``max_bytes`` of it, and return its tag and + sha256; None when nothing is staged there.""" + out.seek(0) + out.truncate() + digest, size = hashlib.sha256(), 0 + async with client.stream("GET", url, headers={"accept-encoding": "identity"}) as response: + if response.status_code == 404: + return None + response.raise_for_status() + async for chunk in response.aiter_raw(_CHUNK_BYTES): + size += len(chunk) + if max_bytes is not None and size > max_bytes: + raise StagingError("the agent staged an object larger than its grant allows") + digest.update(chunk) + await asyncio.to_thread(out.write, chunk) + await asyncio.to_thread(out.flush) + return response.headers.get("etag", ""), digest.hexdigest() + + +def _replace(store: ObjectStore, staged: _Staged, file: Path, digest: str) -> None: + """Put ``file`` over the stored copy of ``staged`` unless the two hold the same bytes. + + The store overwrites only from memory, so replacements go one at a time: an agent rewrites a drained + increment with new bytes only on purpose, and at most one such increment is held at once.""" + stored = hashlib.sha256() + with store.open(staged.object_url) as current: + while chunk := current.read(_CHUNK_BYTES): + stored.update(chunk) + if stored.hexdigest() != digest: + with _REPLACING: + store.put(store.get_object_key(staged.object_url), file.read_bytes(), staged.media_type, allow_overwrite=True) + + +async def _chunks(source: BinaryIO) -> AsyncIterator[bytes]: + try: + while chunk := await asyncio.to_thread(source.read, _CHUNK_BYTES): + yield chunk + finally: + await asyncio.to_thread(source.close) + + +async def _retrying(attempt): + for tried in range(1, _ATTEMPTS + 1): + try: + return await attempt() + except (httpx.TransportError, httpx.HTTPStatusError) as exc: + transient = isinstance(exc, httpx.TransportError) or exc.response.status_code >= 500 + if tried == _ATTEMPTS or not transient: + raise + await asyncio.sleep(tried) + + +async def _each(items, act) -> None: + """``act`` on each of ``items``, a few at a time; the first failure stops the rest.""" + gate = asyncio.Semaphore(_CONCURRENCY) + + async def one(item) -> None: + async with gate: + await act(item) + + tasks = [asyncio.ensure_future(one(item)) for item in items] + try: + await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + +@asynccontextmanager +async def _client(doing: str) -> AsyncIterator[httpx.AsyncClient]: + """A client for staging requests, which are logged by origin alone and fail as a StagingError that says + what was being done.""" + with redacting_request_urls(): + try: + async with httpx.AsyncClient(timeout=_TIMEOUT, follow_redirects=False) as client: + yield client + except httpx.HTTPStatusError as exc: + status = exc.response.status_code + why = {413: "its staging is full", 507: "its disk is full"}.get(status, f"it answered {status}") + raise StagingError(f"{doing} failed: {why}") from None + except httpx.HTTPError: + raise StagingError(f"{doing} failed: the agent could not be reached") from None diff --git a/src/agent_env/a2a_agent/validator.py b/src/agent_env/a2a_agent/validator.py index 84133b62..f2717cb0 100644 --- a/src/agent_env/a2a_agent/validator.py +++ b/src/agent_env/a2a_agent/validator.py @@ -28,6 +28,7 @@ check_changelog_applied, invoke_transfer, ) +from agent_env.a2a_agent.staging import transfer_store from agent_env.config import get_config from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider @@ -661,11 +662,14 @@ def record(*, supported, advertised, save_ok, apply_ok, roundtrip_ok, note=""): A2AAgent.find_extension(apply_agent.a2a_card, A2AAgent.EXT_SNAPSHOT), A2AAgent.SNAPSHOT_METHOD_APPLY_CHANGELOG) a2a_url = apply_agent.a2a_url or apply_agent.api_url + store = transfer_store( + get_config().get_object_store(), a2a_url, apply_agent.a2a_card, sandbox_type=apply_agent.sandbox_type + ) try: call = await asyncio.to_thread( changelog_apply_call, method, - get_config().get_object_store(), + store, agent_name=apply_agent_name, source_url=capture_source, sandbox_type=apply_agent.sandbox_type, @@ -687,6 +691,7 @@ def record(*, supported, advertised, save_ok, apply_ok, roundtrip_ok, note=""): operation="changelog apply", timeout=TRANSFER_TIMEOUT_SECONDS, response_model=ObjectChangelogApplyResponse, + store=store, ) check_changelog_applied(answer, call, agent_name=apply_agent_name) except Exception as e: # noqa: BLE001 diff --git a/src/agent_env/task/teardown.py b/src/agent_env/task/teardown.py index 263de303..4e6ea006 100644 --- a/src/agent_env/task/teardown.py +++ b/src/agent_env/task/teardown.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from typing import NamedTuple, Optional +from agent_env.a2a_agent.staging import drain_all, staged_changelogs from agent_env.env.env import DeployedEnv from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, remove_local_work_dir from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider @@ -82,6 +83,13 @@ async def teardown_run(context: TaskStepContext, *, timeout: float = TERMINATE_T tearing_down() pending = _pending(context) terminated, failed = [], [] + try: # a changelog staged on an agent goes into the store before the agent does + async with asyncio.timeout(timeout): + await drain_all(staged_changelogs(context.metadata)) + except TimeoutError: + logger.warning("teardown: staged changelog increments were not all drained within %gs", timeout) + except asyncio.CancelledError: + return TeardownReport(left=pending) async def take_down(sandbox: RecordedSandbox) -> None: if sandbox not in pending: # torn down already: only a local one's folder can be left diff --git a/src/agent_env/task_step/snapshot_utils/agent_state_capture.py b/src/agent_env/task_step/snapshot_utils/agent_state_capture.py index ff4ce216..fa8c4182 100644 --- a/src/agent_env/task_step/snapshot_utils/agent_state_capture.py +++ b/src/agent_env/task_step/snapshot_utils/agent_state_capture.py @@ -27,6 +27,7 @@ snapshot_save_call, trajectory_mode, ) +from agent_env.a2a_agent.staging import transfer_store from agent_env.config import get_config from agent_env.store.ids import key_segment @@ -82,7 +83,7 @@ async def capture_workspace( capture_key_prefix = ( f"{config.get_artifact_key_prefix()}agent_snapshots/{key_segment(artifact_id)}/{snapshot_version}-{uuid.uuid4().hex[:8]}/" ) - store = config.get_object_store() + store = transfer_store(config.get_object_store(), a2a_url, a2a_card, sandbox_type=sandbox_type) capture_prefix = store.object_url(capture_key_prefix) call = await asyncio.to_thread( @@ -101,6 +102,7 @@ async def capture_workspace( operation="snapshot save", timeout=timeout_seconds, response_model=ObjectSnapshotSaveResponse, + store=store, ) # The agent's answer is only a claim: a snapshot missing either object cannot be restored. @@ -174,7 +176,7 @@ async def read_partial_trajectory( # none may be task_id-only, and would 400 on every tick. if "request" not in get_method: return TrajectoryCapture(reason="trajectory_context_unsupported") - store = get_config().get_object_store() + store = transfer_store(get_config().get_object_store(), a2a_url, a2a_card, sandbox_type=sandbox_type) mode = trajectory_mode(get_method, store, by="context_id", sandbox_type=sandbox_type) if mode is None: return TrajectoryCapture(reason="trajectory_context_unsupported") @@ -189,7 +191,7 @@ async def read_partial_trajectory( return TrajectoryCapture(reason="trajectory_grant_unavailable") try: fetched = await fetch_trajectory( - a2a_url + get_path, {"context_id": context_id}, upload=upload, timeout=timeout_seconds + a2a_url + get_path, {"context_id": context_id}, upload=upload, timeout=timeout_seconds, store=store ) except httpx.HTTPStatusError as exc: if exc.response.status_code == 404: diff --git a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_skill_config.py b/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_skill_config.py index 7b9e4dc6..4054cd05 100644 --- a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_skill_config.py +++ b/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_skill_config.py @@ -16,6 +16,7 @@ parse_response, skill_add_call, ) +from agent_env.a2a_agent.staging import transfer_store from agent_env.config import get_config from agent_env.entity_refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -82,13 +83,16 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: results = rubric_result.get("results", []) inline_passed = any(r.get("id") == "secret_code_inline" and r.get("score") == 1.0 for r in results) add_method, add_path = A2AAgent.operation(skill_ext, "add") - store = get_config().get_object_store() a2a_url = deployed_agent.a2a_url or deployed_agent.api_url async def probe_bundle(name: str, description: str, object_url: str | None) -> bool: """Whether the agent registers an object-backed skill sent as a bundle.""" if object_url is None: return False + store = transfer_store( + get_config().get_object_store(), a2a_url, deployed_agent.a2a_card, + sandbox_type=deployed_agent.sandbox_type, + ) try: call = await asyncio.to_thread( skill_add_call, @@ -112,6 +116,7 @@ async def probe_bundle(name: str, description: str, object_url: str | None) -> b verb="POST", operation="skill add (bundle)", timeout=TRANSFER_TIMEOUT_SECONDS, + store=store, ) result = parse_response(SkillAddResponse, answer, operation="skill add (bundle)") if result.name != name: diff --git a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_trajectory.py b/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_trajectory.py index ca1d7ee8..00b5991f 100644 --- a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_trajectory.py +++ b/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_trajectory.py @@ -12,6 +12,7 @@ fetch_trajectory, trajectory_mode, ) +from agent_env.a2a_agent.staging import transfer_store from agent_env.task_step.snapshot_utils.agent_state_capture import trajectory_object_url from agent_env.entity_refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -90,13 +91,16 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: store = config.get_object_store() probe_prefix = f"{config.get_artifact_key_prefix()}a2a_validator_trajectories/{self.a2a_agent_id}/" sandbox_type = deployed_agent.sandbox_type - if trajectory_mode(get_method, store, by="task_id", sandbox_type=sandbox_type) == "objects": + granting = transfer_store( + store, deployed_agent.a2a_url or deployed_agent.api_url, deployed_agent.a2a_card, sandbox_type=sandbox_type + ) + if trajectory_mode(get_method, granting, by="task_id", sandbox_type=sandbox_type) == "objects": try: - prefix = store.object_url(probe_prefix) + prefix = granting.object_url(probe_prefix) upload = await asyncio.to_thread( - TrajectoryUpload.to, store, trajectory_object_url(prefix, name=task_id, store=store) + TrajectoryUpload.to, granting, trajectory_object_url(prefix, name=task_id, store=granting) ) - await fetch_trajectory(endpoint, {"task_id": task_id}, upload=upload) + await fetch_trajectory(endpoint, {"task_id": task_id}, upload=upload, store=granting) objects_ok = True logger.info("Object trajectory: OK") except Exception as e: diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index 21d7d11e..47943f01 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -22,6 +22,7 @@ invoke_transfer, snapshot_load_call, ) +from agent_env.a2a_agent.staging import StagedObjectStore, transfer_store from agent_env.config import get_config from agent_env.env.env import DeployedEnv, DeployedSandboxEnv, Env from agent_env.providers.sandbox_providers.sandbox_provider import reachable_url, sandbox_request_headers_for_url @@ -230,6 +231,17 @@ def from_dict(cls, data: dict) -> DeployAgentTaskStep: ) async def execute(self, context: TaskStepContext) -> TaskStepContext: + # An agent is recorded in the context only once it is configured, so one that fails on the way there + # (a snapshot or changelog that can't be loaded, say) is closed here, or no teardown would find it. + deployed: list = [] + try: + return await self._execute(context, deployed) + except BaseException: + for agent in deployed: + await agent.close() + raise + + async def _execute(self, context: TaskStepContext, deployed_agents: list) -> TaskStepContext: from agent_env.a2a_agent import A2AAgent from agent_env.a2a_agent.a2a_agent import DEFAULT_A2A_PORT from agent_env.providers.sandbox_providers.sandbox_provider import ( @@ -314,6 +326,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if linked_sandbox is not None: deploy_kwargs["sandbox"] = linked_sandbox deployed = await agent.deploy(**deploy_kwargs) + deployed_agents.append(agent) a2a_url = deployed.a2a_url card = deployed.agent_card logger.info(f"A2A agent '{a2a_agent_id}' deployed at {a2a_url}") @@ -467,7 +480,7 @@ async def _configure_agent_changelog( instance_id = (context.instance_id or "noinstance").replace("/", "_") agent_name = self.agent_name.replace("/", "_") config = get_config() - store = config.get_object_store() + store = transfer_store(config.get_object_store(), a2a_url, card, sandbox_type=sandbox_type) namespace_url = store.object_url(f"{config.get_artifact_key_prefix()}agent_changelog/{instance_id}/{agent_name}") call = await asyncio.to_thread( changelog_enable_call, @@ -485,13 +498,19 @@ async def _configure_agent_changelog( operation="changelog enable", timeout=REPLY_TIMEOUT_SECONDS, response_model=NamespaceChangelogEnableResponse, + store=store, ) - context.metadata.setdefault("agent_changelog", []).append({ + entry = { "agent_name": self.agent_name, "roots": answer.roots, "transfer_mode": call.mode, "object_url": namespace_url, - }) + } + if isinstance(store, StagedObjectStore): + # The increments wait on the agent until they are drained into the store. + entry["staging_url"] = store.namespaces[-1].staging_url + entry["staging_max_object_bytes"] = store.namespaces[-1].max_object_bytes + context.metadata.setdefault("agent_changelog", []).append(entry) logger.info(f"agent-changelog capture enabled on '{self.agent_name}': {namespace_url}") async def _apply_agent_changelog( @@ -507,7 +526,7 @@ async def _apply_agent_changelog( f"Agent '{self.agent_name}' does not support the snapshot " f"'{A2AAgent.SNAPSHOT_METHOD_APPLY_CHANGELOG}' method; cannot apply agent changelog" ) - store = get_config().get_object_store() + store = transfer_store(get_config().get_object_store(), a2a_url, card, sandbox_type=sandbox_type) source_url = self.agent_changelog_object_url call = await asyncio.to_thread( changelog_apply_call, @@ -527,6 +546,7 @@ async def _apply_agent_changelog( operation="changelog apply", timeout=TRANSFER_TIMEOUT_SECONDS, response_model=ObjectChangelogApplyResponse, + store=store, ) check_changelog_applied(answer, call, agent_name=self.agent_name) context_id = answer.context_id @@ -571,10 +591,11 @@ async def _load_snapshot( "snapshot universes must be created via put_existing or put_bundled" ) load_method, load_path = A2AAgent.operation(snapshot_ext, "load") + store = transfer_store(get_config().get_object_store(), a2a_url, card, sandbox_type=sandbox_type) call = await asyncio.to_thread( snapshot_load_call, load_method, - get_config().get_object_store(), + store, agent_name=self.agent_name, bundle_url=universe.bundle_object_url, file_names=set((universe.file_artifact_refs or universe.file_artifact_ids or {}).keys()), @@ -588,6 +609,7 @@ async def _load_snapshot( operation="snapshot load", timeout=TRANSFER_TIMEOUT_SECONDS, response_model=ObjectSnapshotLoadResponse, + store=store, ) loaded_context_id = answer.context_id logger.info( diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index f99e82c2..aa8b9b8f 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -19,6 +19,7 @@ fetch_trajectory, trajectory_mode, ) +from agent_env.a2a_agent.staging import draining, staged_changelogs, transfer_store from agent_env.config.model import MODEL_PARAMS_RESERVED from agent_env.env.gateway.constants import EXT_CLOCK_URI, EXT_TRIGGERS_URI, TRIGGER_IN_FLIGHT_STATUSES from agent_env.store import DuplicateKeyError, get_config @@ -334,6 +335,11 @@ def _warn_on_snapshot_budget(self) -> None: ) async def execute(self, context: TaskStepContext) -> TaskStepContext: + # A changelog staged on the agent moves into the store while it works, and the rest when it is done. + async with draining(staged_changelogs(context.metadata, agent_name=self.agent_name)): + return await self._execute(context) + + async def _execute(self, context: TaskStepContext) -> TaskStepContext: # Minted here, not in `_execute_conversation`, so the series can address the # same session before the first turn. a2a_context_id = self.context_id or uuid.uuid4().hex @@ -512,7 +518,7 @@ async def _execute_conversation( try: turn_traj_uri = await self._fetch_trajectory( target_url, traj_ext_cached, sent_task_id, trajectory_output_prefix, - target_a2a_task_id, sandbox_type=agent.sandbox_type, + target_a2a_task_id, card=card, sandbox_type=agent.sandbox_type, ) except Exception as e: logger.warning(f"Turn {turn+1} trajectory fetch failed (continuing): {e}") @@ -702,10 +708,10 @@ async def _execute_conversation( async def _fetch_trajectory( self, a2a_url: str, traj_ext: dict, a2a_server_task_id: str, trajectory_output_prefix: str, - target_a2a_task_id: str, *, sandbox_type: str | None, + target_a2a_task_id: str, *, sandbox_type: str | None, card: dict | None = None, ) -> str | None: get_method, get_path = A2AAgent.operation(traj_ext, "get") - store = get_config().get_object_store() + store = transfer_store(get_config().get_object_store(), a2a_url, card, sandbox_type=sandbox_type) mode = trajectory_mode(get_method, store, by="task_id", sandbox_type=sandbox_type) if mode is None: raise RuntimeError( @@ -720,7 +726,7 @@ async def _fetch_trajectory( TrajectoryUpload.to, store, trajectory_object_url(trajectory_output_prefix, name=target_a2a_task_id, store=store), ) - fetched = await fetch_trajectory(a2a_url + get_path, {"task_id": a2a_server_task_id}, upload=upload) + fetched = await fetch_trajectory(a2a_url + get_path, {"task_id": a2a_server_task_id}, upload=upload, store=store) return await asyncio.to_thread(store_trajectory, fetched, trajectory_output_prefix, name=target_a2a_task_id) async def _read_env_triggers(self, context: TaskStepContext) -> dict: diff --git a/src/agent_env/task_step/task_steps/teardown_sandboxes.py b/src/agent_env/task_step/task_steps/teardown_sandboxes.py index 2da713b7..7f8f1539 100644 --- a/src/agent_env/task_step/task_steps/teardown_sandboxes.py +++ b/src/agent_env/task_step/task_steps/teardown_sandboxes.py @@ -14,6 +14,7 @@ import logging from typing import Optional +from agent_env.a2a_agent.staging import LAST_DRAIN_SECONDS, drain_all, staged_changelogs from agent_env.entity_refs import EntityRef from agent_env.env.env import DeployedSandboxEnv from agent_env.providers.sandbox_providers.sandbox_provider import ( @@ -84,6 +85,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: logger.warning(f"teardown_sandboxes: not deployed in this run, skipping: {', '.join(missing)}") already = set(context.metadata.get(TORN_DOWN_KEY) or []) pending = {sid: (stype, default) for sid, (stype, default) in targets.items() if sid not in already} + for name in self.agent_names: # a changelog staged on an agent goes into the store before the agent does + await drain_all(staged_changelogs(context.metadata, agent_name=name), within=LAST_DRAIN_SECONDS) results = await asyncio.gather( *(_terminate(sid, stype, default) for sid, (stype, default) in pending.items()), return_exceptions=True, 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 07c1e1ae..56377509 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 @@ -15,6 +15,7 @@ fetch_trajectory, trajectory_mode, ) +from agent_env.a2a_agent.staging import transfer_store from agent_env.config import get_config from agent_env.config.model import ModelParam from agent_env.task_step.context import TaskStepContext @@ -1085,7 +1086,7 @@ async def _fetch_judge_trajectory( return None get_method, get_path = A2AAgent.operation(traj_ext, "get") config = get_config() - store = config.get_object_store() + store = transfer_store(config.get_object_store(), judge_a2a_url, judge_agent_card, sandbox_type=sandbox_type) prefix = store.object_url(f"{config.get_artifact_key_prefix()}judge_trajectories/verifier_id={self.verifier_id}/") mode = trajectory_mode(get_method, store, by="task_id", sandbox_type=sandbox_type) if mode is None: @@ -1110,7 +1111,7 @@ async def _fetch_judge_trajectory( for attempt in range(1, self.DEFAULT_MAX_RETRIES + 1): try: fetched = await fetch_trajectory( - judge_a2a_url + get_path, {"task_id": a2a_server_task_id}, upload=upload + judge_a2a_url + get_path, {"task_id": a2a_server_task_id}, upload=upload, store=store ) return await asyncio.to_thread(store_trajectory, fetched, prefix) except Exception as exc: diff --git a/tst/unit/a2a_agent/staging_test.py b/tst/unit/a2a_agent/staging_test.py new file mode 100644 index 00000000..e2af9fd4 --- /dev/null +++ b/tst/unit/a2a_agent/staging_test.py @@ -0,0 +1,374 @@ +"""An agent the object store's grants cannot reach gets its objects through its own staging routes: pushed +before a call, pulled after it, and a changelog drained into the store while the agent works and before it +goes. The agent here is an SDK agent served in-process at an HTTPS URL, so every request agent-env and the +agent make, the agent's own calls to its staging included, runs without a network.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import httpx +import pytest +from agentenv_protocol.a2a_agent import ( + SKILL_CONFIG_V1, + SNAPSHOT_V1, + TRAJECTORY_V1, + AgentEnvAgent, + AgentIdentity, + BundleSkillRequest, + InlineSkillRequest, + NamespaceChangelogEnableRequest, + NamespaceUploader, + ObjectChangelogApplyRequest, + ObjectSnapshotLoadRequest, + ObjectSnapshotSaveRequest, + TaskObjectTrajectoryRequest, + TaskRequest, + TaskResult, + TaskTrajectoryRequest, + a2a_agent, + download, + extension, + upload, +) +from starlette.testclient import TestClient + +import agent_env.a2a_agent.staging as staging +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.a2a_agent.object_transfer import ( + changelog_apply_call, + changelog_enable_call, + fetch_trajectory, + invoke_transfer, + skill_add_call, + snapshot_load_call, + snapshot_save_call, + TrajectoryUpload, +) +from agent_env.a2a_agent.staging import StagedObjectStore, StagingError, transfer_store +from agent_env.config import reset_config, set_object_store +from agent_env.store import LocalFilesystemObjectStore +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext + +URL = "https://agent.example.test" + + +@a2a_agent(identity=AgentIdentity(name="staging-test", description="test", version="1")) +class _Agent(AgentEnvAgent): + """Moves every object through the SDK's transfer helpers, as any SDK agent does.""" + + workdir: Path + received: dict + uploader: NamespaceUploader | None = None + + async def run(self, request: TaskRequest) -> TaskResult: + return TaskResult.text("ok") + + @extension(SKILL_CONFIG_V1.add) + async def add(self, request: InlineSkillRequest | BundleSkillRequest): + for file in request.skill_bundle.files: + await download(file.object, self.workdir / "skills" / request.name / file.path) + self.received[f"skill/{file.path}"] = (self.workdir / "skills" / request.name / file.path).read_bytes() + return {"name": request.name} + + @extension(TRAJECTORY_V1.get) + async def trajectory(self, request: TaskTrajectoryRequest | TaskObjectTrajectoryRequest): + uploaded = await upload(request.objects.trajectory, b'[{"type": "result"}]') + return {"objects": {"trajectory": uploaded.model_dump(exclude_none=True)}} + + @extension(SNAPSHOT_V1.save) + async def save(self, request: ObjectSnapshotSaveRequest): + trajectory = await upload(request.objects.trajectory, b"conversation") + workspace = await upload(request.objects.workspace, b"w" * 300_000) + return { + "context_id": request.context_id, + "objects": {"trajectory": trajectory.model_dump(), "workspace": workspace.model_dump()}, + } + + @extension(SNAPSHOT_V1.load) + async def load(self, request: ObjectSnapshotLoadRequest): + for name in ("trajectory", "workspace"): + await download(getattr(request.objects, name), self.workdir / "loaded" / name) + self.received[f"snapshot/{name}"] = (self.workdir / "loaded" / name).read_bytes() + return {"context_id": "restored"} + + @extension(SNAPSHOT_V1.changelog.enable) + async def enable_changelog(self, request: NamespaceChangelogEnableRequest): + type(self).uploader = NamespaceUploader(request.write_namespace) + return {"roots": ["/app"]} + + @extension(SNAPSHOT_V1.changelog.apply) + async def apply_changelog(self, request: ObjectChangelogApplyRequest): + for increment in request.increments: + await download(increment.object, self.workdir / "applied" / str(increment.sequence)) + self.received[f"applied/{increment.sequence}"] = (self.workdir / "applied" / str(increment.sequence)).read_bytes() + return {"count": len(request.increments), "context_id": request.target_context_id} + + +@pytest.fixture +def agent(tmp_path, monkeypatch): + """The agent's app, which every HTTP client in this process reaches at ``URL``.""" + monkeypatch.setenv("AGENTENV_STAGING_DIR", str(tmp_path / "agent-staging")) + _Agent.workdir = tmp_path / "agent" + _Agent.received = {} + _Agent.uploader = None + app = _Agent().create_app() + real_async_client = httpx.AsyncClient + monkeypatch.setattr( + httpx, "AsyncClient", lambda **kwargs: real_async_client(transport=httpx.ASGITransport(app=app), **kwargs) + ) + monkeypatch.setattr(httpx, "Client", lambda **kwargs: TestClient(app, base_url=URL, follow_redirects=False)) + return app + + +@pytest.fixture +def store(tmp_path): + store = LocalFilesystemObjectStore(str(tmp_path / "store")) + set_object_store(store) + yield store + reset_config() + + +def _card(app) -> dict: + return TestClient(app).get("/.well-known/agent-card.json").json() + + +def _staged_paths(tmp_path) -> list[str]: + objects = tmp_path / "agent-staging" / "objects" + return sorted(p.relative_to(objects).as_posix() for p in objects.rglob("*") if p.is_file()) if objects.exists() else [] + + +def _method(card, uri, name): + return A2AAgent.operation(A2AAgent.find_extension(card, uri), name) + + +@pytest.mark.asyncio +async def test_a_remote_agent_gets_a_skill_bundle_through_its_staging(agent, store, tmp_path): + store.put("skills/s/SKILL.md", b"# s") + store.put("skills/s/ref/big.bin", b"b" * 3_000_000) + card = _card(agent) + granting = transfer_store(store, URL, card, sandbox_type="modal") + method, path = _method(card, SKILL_CONFIG_V1.uri, "add") + + call = skill_add_call( + method, granting, name="s", description="d", object_url=store.object_url("skills/s"), sandbox_type="modal" + ) + await invoke_transfer(URL + path, call, verb="POST", operation="skill add", timeout=60, store=granting) + + assert isinstance(granting, StagedObjectStore) and call.mode == "objects" + assert _Agent.received == {"skill/SKILL.md": b"# s", "skill/ref/big.bin": b"b" * 3_000_000} + assert _staged_paths(tmp_path) == [] # the call's staging is cleared once it is answered + + +@pytest.mark.asyncio +async def test_what_the_agent_writes_is_in_the_store_when_the_call_returns(agent, store, tmp_path): + card = _card(agent) + granting = transfer_store(store, URL, card, sandbox_type="modal") + save_method, path = _method(card, SNAPSHOT_V1.uri, "save") + prefix = store.object_url("snapshots/one/") + + call = snapshot_save_call( + save_method, granting, agent_name="a", context_id="ctx", capture_prefix=prefix, sandbox_type="modal" + ) + await invoke_transfer(URL + path, call, verb="POST", operation="snapshot save", timeout=60, store=granting) + + assert store.get(store.object_url("snapshots/one/trajectory")) == b"conversation" + assert store.get(store.object_url("snapshots/one/workspace")) == b"w" * 300_000 + assert _staged_paths(tmp_path) == [] + + load_method, path = _method(card, SNAPSHOT_V1.uri, "load") + granting = transfer_store(store, URL, card, sandbox_type="modal") + call = snapshot_load_call( + load_method, granting, agent_name="a", bundle_url=prefix, file_names={"trajectory", "workspace"}, + target_context_id=None, sandbox_type="modal", + ) + await invoke_transfer(URL + path, call, verb="PUT", operation="snapshot load", timeout=60, store=granting) + assert _Agent.received == {"snapshot/trajectory": b"conversation", "snapshot/workspace": b"w" * 300_000} + + +@pytest.mark.asyncio +async def test_a_trajectory_upload_lands_in_the_store(agent, store): + card = _card(agent) + granting = transfer_store(store, URL, card, sandbox_type="modal") + _, path = _method(card, TRAJECTORY_V1.uri, "get") + target = store.object_url("trajectories/t.json") + + fetched = await fetch_trajectory( + URL + path, {"task_id": "t"}, upload=TrajectoryUpload.to(granting, target), store=granting + ) + + assert fetched.object_url == target + assert store.get(target) == b'[{"type": "result"}]' + + +async def _enable_changelog(agent, store) -> tuple[str, staging.StagedNamespace]: + card = _card(agent) + granting = transfer_store(store, URL, card, sandbox_type="modal") + method, path = _method(card, SNAPSHOT_V1.uri, "enable-changelog") + namespace_url = store.object_url("agent_changelog/run/a") + call = changelog_enable_call( + method, granting, agent_name="a", namespace_url=namespace_url, expires_in=3600, sandbox_type="modal" + ) + await invoke_transfer(URL + path, call, verb="POST", operation="changelog enable", timeout=60, store=granting) + [namespace] = granting.namespaces + return namespace_url, namespace + + +@pytest.mark.asyncio +async def test_a_staged_changelog_is_drained_into_the_store_and_applied_to_a_fresh_agent(agent, store, tmp_path): + namespace_url, namespace = await _enable_changelog(agent, store) + await _Agent.uploader.upload("000000.tar", b"first") + await _Agent.uploader.upload("000003.tar", b"second") + + assert await staging.drain(namespace, store) == 2 + assert store.get(f"{namespace_url}/000000.tar") == b"first" + assert store.get(f"{namespace_url}/000003.tar") == b"second" + assert _staged_paths(tmp_path) == [] + assert await staging.drain(namespace, store) == 0 + + card = _card(agent) + granting = transfer_store(store, URL, card, sandbox_type="modal") + method, path = _method(card, SNAPSHOT_V1.uri, "apply-changelog") + call = changelog_apply_call(method, granting, agent_name="a", source_url=namespace_url, sandbox_type="modal") + await invoke_transfer(URL + path, call, verb="PUT", operation="changelog apply", timeout=60, store=granting) + assert _Agent.received == {"applied/0": b"first", "applied/3": b"second"} + + +@pytest.mark.asyncio +async def test_a_prompt_drains_as_the_agent_works_and_once_more_when_it_ends(agent, store, monkeypatch): + monkeypatch.setattr(staging, "DRAIN_INTERVAL_SECONDS", 0.05) + namespace_url, namespace = await _enable_changelog(agent, store) + + async with staging.draining([namespace]): + await _Agent.uploader.upload("000000.tar", b"mid-prompt") + for _ in range(100): + if store.get_object_metadata_at(f"{namespace_url}/000000.tar") is not None: + break + await asyncio.sleep(0.02) + assert store.get(f"{namespace_url}/000000.tar") == b"mid-prompt" # before the prompt ends + await _Agent.uploader.upload("000001.tar", b"last") + assert store.get(f"{namespace_url}/000001.tar") == b"last" + + +@pytest.mark.asyncio +async def test_a_run_torn_down_drains_its_staged_changelogs_first(agent, store): + namespace_url, namespace = await _enable_changelog(agent, store) + await _Agent.uploader.upload("000000.tar", b"before teardown") + context = TaskStepContext(metadata={"agent_changelog": [ + {"agent_name": "a", "object_url": namespace_url, "transfer_mode": "objects", "staging_url": namespace.staging_url}, + ]}) + + await teardown_run(context) + + assert store.get(f"{namespace_url}/000000.tar") == b"before teardown" + + +def test_staging_is_only_for_agents_the_grants_cannot_reach(agent, store): + card = _card(agent) + assert transfer_store(store, URL, card, sandbox_type="local") is store # the local grant server reaches it + assert transfer_store(store, URL, {}, sandbox_type="modal") is store # no staging: the inline forms, or a refusal + assert transfer_store(store, "http://agent.example.test", card, sandbox_type="modal") is store # no HTTPS grant + assert isinstance(transfer_store(store, URL, card, sandbox_type="modal"), StagedObjectStore) + + +@pytest.mark.asyncio +async def test_a_full_staging_fails_the_call_without_naming_a_staged_path(tmp_path, monkeypatch, store): + monkeypatch.setenv("AGENTENV_STAGING_MAX_BYTES", "1000") + app = _Agent().create_app() + real_async_client = httpx.AsyncClient + monkeypatch.setattr( + httpx, "AsyncClient", lambda **kwargs: real_async_client(transport=httpx.ASGITransport(app=app), **kwargs) + ) + _Agent.workdir, _Agent.received = tmp_path / "agent", {} + store.put("skills/s/SKILL.md", b"x" * 5000) + card = _card(app) + granting = transfer_store(store, URL, card, sandbox_type="modal") + method, path = _method(card, SKILL_CONFIG_V1.uri, "add") + call = skill_add_call( + method, granting, name="s", description="d", object_url=store.object_url("skills/s"), sandbox_type="modal" + ) + + with pytest.raises(StagingError, match="staging an object on the agent failed: its staging is full") as raised: + await invoke_transfer(URL + path, call, verb="POST", operation="skill add", timeout=60, store=granting) + assert granting.endpoint not in str(raised.value) + assert granting.max_bytes == 1000 # refused from the card's limit, before anything was sent + + +class _CountingStore(LocalFilesystemObjectStore): + lookups = 0 + + def get_object_metadata_at(self, object_url): + self.lookups += 1 + return super().get_object_metadata_at(object_url) + + +@pytest.mark.asyncio +async def test_a_push_sends_the_objects_as_the_call_described_them_without_asking_the_store_again(agent, tmp_path): + store = _CountingStore(str(tmp_path / "counted")) + set_object_store(store) + try: + files = {"SKILL.md": b"# s", "a.txt": b"a" * 10, "b.txt": b"b" * 20} + for name, data in files.items(): + store.put(f"skills/s/{name}", data) + card = _card(agent) + granting = transfer_store(store, URL, card, sandbox_type="modal") + method, path = _method(card, SKILL_CONFIG_V1.uri, "add") + call = skill_add_call( + method, granting, name="s", description="d", object_url=store.object_url("skills/s"), sandbox_type="modal" + ) + described = store.lookups + + await invoke_transfer(URL + path, call, verb="POST", operation="skill add", timeout=60, store=granting) + + assert store.lookups == described + assert _Agent.received == {f"skill/{name}": data for name, data in files.items()} + finally: + reset_config() + + +@pytest.mark.asyncio +async def test_a_pull_stops_at_the_size_its_grant_allows(agent, store): + granting = transfer_store(store, URL, _card(agent), sandbox_type="modal") + grant = granting.issue_write_grant(store.object_url("out/t.json"), media_type="application/json", max_bytes=10) + async with httpx.AsyncClient() as client: # an agent whose own client ignores the limit + assert (await client.put(grant.url, content=b"x" * 100)).status_code == 201 + + with pytest.raises(StagingError, match="larger than its grant allows"): + await granting.pull() + assert store.get_object_metadata_at(store.object_url("out/t.json")) is None + + +@pytest.mark.asyncio +async def test_an_increment_rewritten_after_it_was_drained_replaces_the_stored_copy(agent, store): + namespace_url, namespace = await _enable_changelog(agent, store) + await _Agent.uploader.upload("000000.tar", b"first") + await staging.drain(namespace, store) + await _Agent.uploader.upload("000000.tar", b"first") # the same bytes again, as a retried upload sends + await staging.drain(namespace, store) + assert store.get(f"{namespace_url}/000000.tar") == b"first" + + await _Agent.uploader.upload("000000.tar", b"second") + await staging.drain(namespace, store) + assert store.get(f"{namespace_url}/000000.tar") == b"second" + + +@pytest.mark.asyncio +async def test_a_drain_that_keeps_failing_is_reported_once(monkeypatch, caplog, tmp_path): + monkeypatch.setattr(staging, "DRAIN_INTERVAL_SECONDS", 0.01) + local = LocalFilesystemObjectStore(str(tmp_path / "store")) + unreachable = staging.StagedNamespace("https://gone.example.test/ext/staging/" + "n" * 32, local.object_url("ns")) + + def refuse(request): + raise httpx.ConnectError("refused", request=request) + + real_async_client = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: real_async_client(transport=httpx.MockTransport(refuse), **kwargs)) + monkeypatch.setattr(staging, "get_config", lambda: type("C", (), {"get_object_store_at": lambda self, url: local})()) + with caplog.at_level("WARNING", logger=staging.__name__): + async with staging.draining([unreachable]): + await asyncio.sleep(0.2) + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 2 # once while the prompt runs, once for the last drain + assert "the agent could not be reached" in warnings[0].getMessage() diff --git a/tst/unit/task_step/deploy_agent_cleanup_test.py b/tst/unit/task_step/deploy_agent_cleanup_test.py new file mode 100644 index 00000000..d887d21e --- /dev/null +++ b/tst/unit/task_step/deploy_agent_cleanup_test.py @@ -0,0 +1,57 @@ +"""An agent that fails to be configured after its sandbox exists is closed, since a run's teardown only finds the +agents its context records, and deploy_agent records one only once it is configured.""" + +from types import SimpleNamespace + +import pytest + +import agent_env.a2a_agent as a2a_agent_package +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep + + +class _Agent: + metadata: dict = {} + + def __init__(self) -> None: + self.closed = 0 + + async def deploy(self, **_): + return SimpleNamespace( + a2a_url="https://agent.example.test", agent_card={}, sandbox_id="sb-1", sandbox_type="fake", + instance_id="i-1", network_policy=None, + ) + + async def close(self) -> None: + self.closed += 1 + + +@pytest.mark.asyncio +async def test_an_agent_whose_snapshot_load_fails_is_closed_and_not_recorded(monkeypatch): + agent = _Agent() + monkeypatch.setattr(a2a_agent_package.A2AAgent, "get", classmethod(lambda cls, *a, **k: agent)) + + async def unloadable(self, *args, **kwargs): + raise RuntimeError("staging an object on the agent failed: the agent could not be reached") + + monkeypatch.setattr(DeployAgentTaskStep, "_load_snapshot", unloadable) + step = DeployAgentTaskStep( + id="deploy", version=None, a2a_agent_id="a", agent_name="solver", agent_snapshot_files_artifact_id="u" + ) + context = TaskStepContext() + + with pytest.raises(RuntimeError, match="could not be reached"): + await step.execute(context) + assert agent.closed == 1 + assert context.deployed_agents == [] + + +@pytest.mark.asyncio +async def test_a_configured_agent_is_recorded_and_left_running(monkeypatch): + agent = _Agent() + monkeypatch.setattr(a2a_agent_package.A2AAgent, "get", classmethod(lambda cls, *a, **k: agent)) + step = DeployAgentTaskStep(id="deploy", version=None, a2a_agent_id="a", agent_name="solver") + context = await step.execute(TaskStepContext()) + + assert agent.closed == 0 + assert [a.sandbox_id for a in context.deployed_agents] == ["sb-1"]