From 5736df1f7aee29a22710d50e0f191efc24c74331 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 7 Oct 2026 09:21:01 -0700 Subject: [PATCH 1/2] fix(a2a): agents reach their own staging over loopback A staged grant names a path on the agent's own public URL, and a sandbox can't always call that URL: on one provider the in-sandbox route to the sandbox's public hostname serves a certificate that doesn't verify there, so every staged transfer failed. Core's staged grants now also carry an AgentEnv-Staging-Path header, the path the agent's own server serves the URL at. The protocol's transfer helpers send such a request to http://127.0.0.1:$A2A_PORT first, and to the grant's URL only when nothing listens there; transfers.loopback_url gives the same URL to agents with transfer code of their own. Agents and cores on either side of the change still work together: a grant without the header, or a helper that ignores it, uses the grant's URL. Co-Authored-By: Claude Opus 5.5 (1M context) --- packages/agentenv-protocol/README.md | 11 +- .../src/agentenv_protocol/transfers.py | 142 ++++++++++++----- .../tests/test_a2a_object_transfers.py | 144 ++++++++++++++++++ src/agent_env/a2a_agent/object_transfer.py | 5 +- src/agent_env/a2a_agent/staging.py | 45 ++++-- tst/unit/a2a_agent/staging_test.py | 81 ++++++++-- 6 files changed, 365 insertions(+), 63 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index fb719636..29f3a2b6 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -622,8 +622,15 @@ itself. Every SDK agent serves the staging extension, `urn:agentenv:staging/v1`, a small object store at `/ext/staging` on its own server. Before a call AgentEnv pushes what the agent will read into it; after the call it pulls what the agent wrote. The grants it sends are the ordinary ones above, with URLs -naming staged paths on the agent's own URL, so handlers and helpers are -unchanged. +naming staged paths on the agent's own URL, so handlers are unchanged. + +A sandbox can't always call its own public URL, so each of those grants also +carries an `AgentEnv-Staging-Path` header: the path the agent's own server +serves the URL at, such as `/ext/staging/{path}`. The helpers send such a +request to that server over loopback, `http://127.0.0.1:$A2A_PORT/ext/staging/{path}` +(AgentEnv deploys every agent with `A2A_PORT` set), and to the grant's URL only +when nothing listens there. An agent with transfer code of its own gets the same +URL from `transfers.loopback_url(url, headers)`. | Route | Does | | --- | --- | diff --git a/packages/agentenv-protocol/src/agentenv_protocol/transfers.py b/packages/agentenv-protocol/src/agentenv_protocol/transfers.py index 65dda9cc..6deca71d 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/transfers.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/transfers.py @@ -10,14 +10,14 @@ import re import io import tempfile -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from contextlib import contextmanager from contextvars import ContextVar from datetime import datetime, timezone from functools import partial from pathlib import Path from typing import Annotated, BinaryIO, Literal, TypeVar -from urllib.parse import urlsplit +from urllib.parse import urlsplit, urlunsplit import httpx from pydantic import ( @@ -51,6 +51,11 @@ _CONNECT_TIMEOUT_SECONDS + TRANSFER_IDLE_TIMEOUT_SECONDS ) + sum(_RETRY_BACKOFF_SECONDS * 2**attempt for attempt in range(TRANSFER_ATTEMPTS - 1)) _T = TypeVar("_T") +# A grant naming a path on its holder's own staging carries this header, whose value is that path on +# the holder's server: a sandbox can't always call its own public URL, so the helpers reach the server +# over loopback, on the port agent-env deploys it on. +STAGING_PATH_HEADER = "AgentEnv-Staging-Path" +_SERVER_PORT_ENV = "A2A_PORT" def _https_url(value: str) -> str: @@ -281,6 +286,51 @@ def _is_s3_request_timeout(response: httpx.Response) -> bool: return False +def loopback_url(url: str, headers: Mapping[str, str] | None) -> str | None: + """Where the holder's own server answers a grant's ``url`` over loopback: when the grant names a path + on the holder's staging (an ``AgentEnv-Staging-Path`` header ending the URL's path) and ``A2A_PORT`` + says which port that server listens on. None otherwise.""" + path = httpx.Headers(headers or {}).get(STAGING_PATH_HEADER) + port = os.environ.get(_SERVER_PORT_ENV, "") + if not path or not path.startswith("/") or not port.isdigit() or not 0 < int(port) < 65536: + return None + parsed = urlsplit(url) + if not parsed.path.endswith(path): + return None + return urlunsplit(("http", f"127.0.0.1:{int(port)}", path, parsed.query, "")) + + +def _destinations(url: str, headers: Mapping[str, str] | None) -> tuple[str, ...]: + """Where a grant's request goes: the holder's own server over loopback first when the grant names a + path on its staging, then the grant's URL, for a server that isn't listening there.""" + loopback = loopback_url(url, headers) + return (url,) if loopback is None else (loopback, url) + + +async def _send_to_first_listening( + urls: tuple[str, ...], send: Callable[[str], Awaitable[_T]] +) -> _T: + """``send`` to the first of ``urls`` whose server takes the connection.""" + *earlier, last = urls + for url in earlier: + try: + return await send(url) + except httpx.ConnectError: + pass + return await send(last) + + +def _send_to_first_listening_sync(urls: tuple[str, ...], send: Callable[[str], _T]) -> _T: + """``_send_to_first_listening`` for a sync client.""" + *earlier, last = urls + for url in earlier: + try: + return send(url) + except httpx.ConnectError: + pass + return send(last) + + def _check_unexpired(expires_at: datetime) -> None: if expires_at <= datetime.now(_UTC): raise TransferError("grant_expired", "The object-transfer grant has expired.") @@ -401,20 +451,21 @@ async def upload(target: WriteObject, source: Path | bytes) -> Uploaded: headers = httpx.Headers(target.write.headers or {}) headers.setdefault("content-type", target.media_type) headers.setdefault("content-length", str(size_bytes)) + urls = _destinations(target.write.url, target.write.headers) async def put() -> Uploaded: - with ( - _SourceStream(source, size_bytes) as stream, - _transfer_request(target.write.expires_at), - ): + with _transfer_request(target.write.expires_at): async with httpx.AsyncClient( follow_redirects=False, timeout=_TRANSFER_TIMEOUT ) as client: - response = await client.put( - target.write.url, headers=headers, content=stream - ) - _raise_for_transfer_status(response) - return stream.uploaded() + + async def to(url: str) -> Uploaded: + with _SourceStream(source, size_bytes) as stream: + response = await client.put(url, headers=headers, content=stream) + _raise_for_transfer_status(response) + return stream.uploaded() + + return await _send_to_first_listening(urls, to) return await _retrying(target.write.expires_at, put) @@ -434,28 +485,35 @@ async def download(source: ReadObject, destination: Path) -> None: temporary = Path(filename) headers = httpx.Headers(source.read.headers or {}) headers["accept-encoding"] = "identity" + urls = _destinations(source.read.url, source.read.headers) async def get() -> tuple[int, str]: size_bytes = 0 digest = hashlib.sha256() with _transfer_request(source.read.expires_at): - async with ( - httpx.AsyncClient( - follow_redirects=False, timeout=_TRANSFER_TIMEOUT - ) as client, - client.stream("GET", source.read.url, headers=headers) as response, - ): - _raise_for_transfer_status(response) - stream = await asyncio.to_thread(temporary.open, "wb") + async with httpx.AsyncClient( + follow_redirects=False, timeout=_TRANSFER_TIMEOUT + ) as client: + response = await _send_to_first_listening( + urls, + lambda url: client.send( + client.build_request("GET", url, headers=headers), stream=True + ), + ) try: - async for chunk in response.aiter_raw(_CHUNK_BYTES): - size_bytes += len(chunk) - if size_bytes > source.max_bytes: - raise TransferError("transfer_too_large", _TOO_LARGE) - digest.update(chunk) - await asyncio.to_thread(stream.write, chunk) + _raise_for_transfer_status(response) + stream = await asyncio.to_thread(temporary.open, "wb") + try: + async for chunk in response.aiter_raw(_CHUNK_BYTES): + size_bytes += len(chunk) + if size_bytes > source.max_bytes: + raise TransferError("transfer_too_large", _TOO_LARGE) + digest.update(chunk) + await asyncio.to_thread(stream.write, chunk) + finally: + await asyncio.to_thread(stream.close) finally: - await asyncio.to_thread(stream.close) + await response.aclose() return size_bytes, digest.hexdigest() try: @@ -485,24 +543,28 @@ def _post_to_namespace( fields = {**grant.fields, grant.path_field: object_path} fields.setdefault("Content-Type", "application/octet-stream") with ( - _SourceStream(source, size_bytes) as stream, _transfer_request(target.expires_at), httpx.Client(follow_redirects=False, timeout=_TRANSFER_TIMEOUT) as client, ): - response = client.post( - grant.url, - data=fields, - files={ - grant.file_field: ( - object_path.rsplit("/", 1)[-1], - stream, - "application/octet-stream", + + def to(url: str) -> Uploaded: + with _SourceStream(source, size_bytes) as stream: + response = client.post( + url, + data=fields, + files={ + grant.file_field: ( + object_path.rsplit("/", 1)[-1], + stream, + "application/octet-stream", + ) + }, + headers=grant.headers, ) - }, - headers=grant.headers, - ) - _raise_for_transfer_status(response) - return stream.uploaded() + _raise_for_transfer_status(response) + return stream.uploaded() + + return _send_to_first_listening_sync(_destinations(grant.url, grant.headers), to) class NamespaceUploader: diff --git a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py index ba10fe4e..e7b1d36a 100644 --- a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py +++ b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py @@ -964,3 +964,147 @@ def handle(request: httpx.Request) -> httpx.Response: assert "transfer-encoding" not in request.headers assert int(request.headers["content-length"]) == len(request.content) assert from_bytes == from_file + + +_PUBLIC = "https://agent.example.test/sandbox/vm-8000" # the provider routes this prefix to the agent +_STAGED = "/ext/staging/" + "c" * 32 + "/0" +_NAMESPACE = "/ext/staging/" + "n" * 32 + + +def _on_own_staging(content: bytes) -> tuple[WriteObject, ReadObject, WriteNamespaceGrant]: + """Grants naming paths on their holder's own staging, as agent-env stages them.""" + headers = {transfers.STAGING_PATH_HEADER: _STAGED} + write = WriteObject( + media_type="application/json", + max_bytes=1024, + write=HttpPutGrant(kind="http-put", url=_PUBLIC + _STAGED, expires_at=_expiry(), headers=headers), + ) + read = ReadObject( + media_type="application/json", + max_bytes=1024, + size_bytes=len(content), + sha256=hashlib.sha256(content).hexdigest(), + read=HttpGetGrant(kind="http-get", url=_PUBLIC + _STAGED, expires_at=_expiry(), headers=headers), + ) + namespace = _namespace_grant().model_copy( + update={ + "write": HttpPostPolicyGrant( + kind="http-post-policy", + url=_PUBLIC + _NAMESPACE, + fields={}, + path_field="key", + file_field="file", + headers={transfers.STAGING_PATH_HEADER: _NAMESPACE}, + ) + } + ) + return write, read, namespace + + +@pytest.mark.asyncio +async def test_a_grant_on_its_holders_own_staging_goes_to_its_server_over_loopback( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("A2A_PORT", "8123") + content = b'{"trajectory":[]}' + sent: list[str] = [] + + def handle(request: httpx.Request) -> httpx.Response: + sent.append(f"{request.method} {request.url}") + if request.url.host != "127.0.0.1": + raise httpx.ConnectError("the sandbox can't call its own public URL", request=request) + request.read() + return _served(content) if request.method == "GET" else httpx.Response(201) + + _route(monkeypatch, handle) + write, read, namespace = _on_own_staging(content) + + await upload(write, content) + await download(read, tmp_path / "trajectory.json") + await NamespaceUploader(namespace).upload("000000.tar", b"increment") + + assert sent == [ + f"PUT http://127.0.0.1:8123{_STAGED}", + f"GET http://127.0.0.1:8123{_STAGED}", + f"POST http://127.0.0.1:8123{_NAMESPACE}", + ] + assert (tmp_path / "trajectory.json").read_bytes() == content + + +@pytest.mark.asyncio +async def test_a_holder_whose_server_isnt_listening_on_loopback_uses_the_grants_url( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("A2A_PORT", "8123") + content = b'{"trajectory":[]}' + sent: list[str] = [] + bodies: list[bytes] = [] + + def handle(request: httpx.Request) -> httpx.Response: + sent.append(request.url.host) + if request.url.host == "127.0.0.1": + raise httpx.ConnectError("refused", request=request) + bodies.append(request.read()) + return _served(content) if request.method == "GET" else httpx.Response(201) + + _route(monkeypatch, handle) + write, read, namespace = _on_own_staging(content) + + uploaded = await upload(write, content) + await download(read, tmp_path / "trajectory.json") + increment = await NamespaceUploader(namespace).upload("000000.tar", b"increment") + + assert sent == ["127.0.0.1", "agent.example.test"] * 3 + assert bodies[0] == content and uploaded.size_bytes == len(content) + assert b"increment" in bodies[2] and increment.size_bytes == len(b"increment") + assert (tmp_path / "trajectory.json").read_bytes() == content + + +@pytest.mark.asyncio +async def test_an_answer_from_the_holders_own_server_is_final(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("A2A_PORT", "8123") + requests: list[httpx.Request] = [] + _route(monkeypatch, _respond(requests, 404)) + write, _, _ = _on_own_staging(b"") + + with pytest.raises(TransferError) as exc_info: + await upload(write, b"{}") + + assert exc_info.value.code == "transfer_rejected" + assert [request.url.host for request in requests] == ["127.0.0.1"] + + +@pytest.mark.parametrize( + ("port", "path"), + [ + (None, _STAGED), # agent-env didn't deploy the holder + ("http", _STAGED), + ("0", _STAGED), + ("65536", _STAGED), + ("8123", None), # not staged + ("8123", "/ext/staging/" + "d" * 32 + "/0"), # not this URL's path + ("8123", _STAGED.lstrip("/")), + ], +) +def test_a_grant_goes_to_its_url_unless_it_names_a_path_its_holders_known_port_serves( + monkeypatch: pytest.MonkeyPatch, port: str | None, path: str | None +) -> None: + if port is None: + monkeypatch.delenv("A2A_PORT", raising=False) + else: + monkeypatch.setenv("A2A_PORT", port) + headers = None if path is None else {transfers.STAGING_PATH_HEADER: path} + + assert transfers.loopback_url(_PUBLIC + _STAGED, headers) is None + + +def test_the_loopback_url_keeps_the_grants_query_and_reads_the_header_in_any_case( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("A2A_PORT", "8123") + headers = {transfers.STAGING_PATH_HEADER.lower(): _STAGED} + + assert ( + transfers.loopback_url(f"{_PUBLIC}{_STAGED}?part=1", headers) + == f"http://127.0.0.1:8123{_STAGED}?part=1" + ) diff --git a/src/agent_env/a2a_agent/object_transfer.py b/src/agent_env/a2a_agent/object_transfer.py index 4dcc2b37..5f1d3156 100644 --- a/src/agent_env/a2a_agent/object_transfer.py +++ b/src/agent_env/a2a_agent/object_transfer.py @@ -60,8 +60,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." + "server, and the agent could not reach them there. An agent on a current SDK reaches its staging over " + "loopback, on its A2A_PORT; an older one goes through its own public URL, which its sandbox provider may " + "not allow. An object store whose grants reach the agent avoids staging." ) # The time budget of one transfer, outermost first: a grant (the issuing store's diff --git a/src/agent_env/a2a_agent/staging.py b/src/agent_env/a2a_agent/staging.py index c7700fb9..00b622d4 100644 --- a/src/agent_env/a2a_agent/staging.py +++ b/src/agent_env/a2a_agent/staging.py @@ -26,8 +26,14 @@ 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 agentenv_protocol.a2a_agent import STAGING_ENDPOINT, STAGING_V1_URI +from agentenv_protocol.transfers import ( + STAGING_PATH_HEADER, + HttpGetGrant, + HttpPostPolicyGrant, + HttpPutGrant, + redacting_request_urls, +) from agent_env.config import get_config from agent_env.store.base import ObjectAlreadyExistsError @@ -89,8 +95,14 @@ def staged_store(store: ObjectStore, a2a_url: str, card: Mapping[str, Any] | Non def _staged(store: ObjectStore, endpoint: str, card: Mapping[str, Any] | None) -> StagedObjectStore: - limit = _staging_params(card).get("max_bytes") - return StagedObjectStore(store, endpoint, max_bytes=limit if isinstance(limit, int) and limit > 0 else None) + params = _staging_params(card) + limit = params.get("max_bytes") + return StagedObjectStore( + store, + endpoint, + endpoint_path=params["endpoint"].rstrip("/"), + max_bytes=limit if isinstance(limit, int) and limit > 0 else None, + ) @dataclass(frozen=True) @@ -116,13 +128,20 @@ class StagedObjectStore: 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. + + Each grant also names its path on the agent's own server, below ``endpoint_path``, in an + ``AgentEnv-Staging-Path`` header: a sandbox can't always call its own public URL, so the agent's + transfer helpers reach the path over loopback instead. """ supports_transfer_grants = True - def __init__(self, store: ObjectStore, endpoint: str, *, max_bytes: int | None = None) -> None: + def __init__( + self, store: ObjectStore, endpoint: str, *, endpoint_path: str = STAGING_ENDPOINT, max_bytes: int | None = None + ) -> None: self.store = store self.endpoint = endpoint + self.endpoint_path = endpoint_path # where the agent's own server serves ``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] = [] @@ -142,20 +161,22 @@ def get_object_metadata_at(self, object_url: str) -> ObjectMetadata | None: 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)) + url = self._url(self._stage(self._reads, object_url).path) + return HttpGetGrant(kind="http-get", url=url, expires_at=self._expiry(expires_in), headers=self._on_agent(url)) 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)) + url = self._url(self._stage(self._writes, object_url, media_type, max_bytes).path) + return HttpPutGrant(kind="http-put", url=url, expires_at=self._expiry(expires_in), headers=self._on_agent(url)) 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"), + write=HttpPostPolicyGrant( + kind="http-post-policy", url=url, fields={}, path_field="key", file_field="file", headers=self._on_agent(url) + ), expires_at=self._expiry(expires_in), ) @@ -200,6 +221,10 @@ def _stage( def _url(self, path: str) -> str: return f"{self.endpoint}/{self._call}" + (f"/{quote(path)}" if path else "") + def _on_agent(self, url: str) -> dict[str, str]: + """The header naming the path the agent's own server serves ``url``, a URL below the endpoint, at.""" + return {STAGING_PATH_HEADER: self.endpoint_path + url[len(self.endpoint):]} + 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) diff --git a/tst/unit/a2a_agent/staging_test.py b/tst/unit/a2a_agent/staging_test.py index 7ac08bd0..8743f050 100644 --- a/tst/unit/a2a_agent/staging_test.py +++ b/tst/unit/a2a_agent/staging_test.py @@ -1,7 +1,9 @@ """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.""" +agent make, the agent's own calls to its staging included, runs without a network. Each flow runs twice: +with the agent reaching its staging through its public URL, and in a sandbox that can't call its own public +URL, where it reaches its staging over loopback.""" from __future__ import annotations @@ -32,6 +34,9 @@ extension, upload, ) +from agentenv_protocol import transfers +from agentenv_protocol.transfers import STAGING_PATH_HEADER, loopback_url +from starlette.responses import PlainTextResponse from starlette.testclient import TestClient import agent_env.a2a_agent.staging as staging @@ -109,19 +114,45 @@ async def apply_changelog(self, request: ObjectChangelogApplyRequest): 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() +def _calling_its_own_url_fails(app): + """``app`` in a sandbox that can't call its own public URL: a request the agent sends its staging there + fails, as one sent over loopback doesn't. Only the agent's own requests carry the staging path header.""" + + async def sandboxed(scope, receive, send): + headers = dict(scope.get("headers") or ()) + host = headers.get(b"host", b"").split(b":")[0] + if scope["type"] == "http" and STAGING_PATH_HEADER.lower().encode() in headers and host != b"127.0.0.1": + await PlainTextResponse("the sandbox can't call its own public URL", status_code=502)(scope, receive, send) + return + await app(scope, receive, send) + + return sandboxed + + +def _serve(app, monkeypatch) -> None: + """Send every HTTP client in this process to ``app``, whatever URL it asks for.""" 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)) + + +@pytest.fixture(params=["public URL", "loopback"]) +def agent(request, tmp_path, monkeypatch): + """The agent's app, which every HTTP client in this process reaches at ``URL``: reaching its staging + through its public URL, or over loopback from a sandbox that can't call that 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() + if request.param == "loopback": + monkeypatch.setenv("A2A_PORT", "8000") + app = _calling_its_own_url_fails(app) + else: + monkeypatch.delenv("A2A_PORT", raising=False) + _serve(app, monkeypatch) return app @@ -302,6 +333,38 @@ def test_staging_is_only_for_agents_the_grants_cannot_reach(agent, store): assert isinstance(transfer_store(store, URL, card, sandbox_type="modal"), StagedObjectStore) +def test_each_staged_grant_names_its_path_on_the_agents_own_server(store, monkeypatch): + """The provider routes ``/sandbox/vm-8000`` on its host to the agent, whose server serves ``/ext/staging``.""" + monkeypatch.setenv("A2A_PORT", "8000") + public = f"{URL}/sandbox/vm-8000" + granting = transfer_store(store, public, _card(_Agent().create_app()), sandbox_type="modal") + read = granting.issue_read_grant(store.put("in/x", b"x")) + write = granting.issue_write_grant(store.object_url("out/y"), media_type="application/json", max_bytes=10) + policy = granting.issue_upload_policy(store.object_url("ns/"), max_object_bytes=10, expires_in=60).write + + for grant in (read, write, policy): + path = grant.headers[STAGING_PATH_HEADER] + assert path.startswith("/ext/staging/") and grant.url == public + path + assert loopback_url(grant.url, grant.headers) == f"http://127.0.0.1:8000{path}" + + +@pytest.mark.asyncio +async def test_an_agent_that_cant_call_its_own_url_and_has_no_loopback_port_fails_saying_why(tmp_path, monkeypatch, store): + monkeypatch.setenv("AGENTENV_STAGING_DIR", str(tmp_path / "agent-staging")) + monkeypatch.delenv("A2A_PORT", raising=False) + monkeypatch.setattr(transfers, "_RETRY_BACKOFF_SECONDS", 0) + app = _calling_its_own_url_fails(_Agent().create_app()) + _serve(app, monkeypatch) + card = _card(app) + granting = transfer_store(store, URL, card, sandbox_type="modal") + _, path = _method(card, TRAJECTORY_V1.uri, "get") + target = store.object_url("trajectories/t.json") + + with pytest.raises(httpx.HTTPStatusError, match="reaches its staging over loopback, on its A2A_PORT"): + await fetch_trajectory(URL + path, {"task_id": "t"}, upload=TrajectoryUpload.to(granting, target), store=granting) + assert store.get_object_metadata_at(target) is None + + @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") From f6e2b99a5f009c690ac692884d9f5a1c270361c5 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 7 Oct 2026 09:54:55 -0700 Subject: [PATCH 2/2] fix(a2a): the SDK's server reaches its staging on the port it took the request on A2A_PORT reaches an agent's server only when whatever starts it sets it: agent-env does for an agent it deploys from an image, but an installed agent's own install commands start it. The SDK's server now serves each request serving_on the port its socket took it on, and loopback_url prefers that port, so an SDK agent reaches its staging over loopback however it was started. A2A_PORT stays the port for transfers made outside a request, as from a hook process. Co-Authored-By: Claude Opus 5.5 (1M context) --- packages/agentenv-protocol/README.md | 11 +++--- .../agentenv_protocol/a2a_agent/framework.py | 23 +++++++++++-- .../src/agentenv_protocol/transfers.py | 21 +++++++++--- .../tests/test_a2a_object_transfers.py | 34 +++++++++++++++++++ src/agent_env/a2a_agent/object_transfer.py | 4 +-- tst/unit/a2a_agent/staging_test.py | 14 ++++---- 6 files changed, 89 insertions(+), 18 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 29f3a2b6..cd84562d 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -627,10 +627,13 @@ naming staged paths on the agent's own URL, so handlers are unchanged. A sandbox can't always call its own public URL, so each of those grants also carries an `AgentEnv-Staging-Path` header: the path the agent's own server serves the URL at, such as `/ext/staging/{path}`. The helpers send such a -request to that server over loopback, `http://127.0.0.1:$A2A_PORT/ext/staging/{path}` -(AgentEnv deploys every agent with `A2A_PORT` set), and to the grant's URL only -when nothing listens there. An agent with transfer code of its own gets the same -URL from `transfers.loopback_url(url, headers)`. +request to that server over loopback, `http://127.0.0.1:{port}/ext/staging/{path}`, +and to the grant's URL only when nothing listens there. The port is the one the +SDK's server took the request being served on, or else `A2A_PORT` (set by AgentEnv +on agents it deploys from an image), as in a separate process such as a hook. +An agent with transfer code of its own gets the same URL from +`transfers.loopback_url(url, headers)`, and an agent with a server of its own +marks each request with `transfers.serving_on(port)`. | Route | Does | | --- | --- | diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py index 500aaad5..5c3c3436 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py @@ -22,9 +22,11 @@ from pydantic_core import core_schema from starlette.applications import Starlette from starlette.exceptions import HTTPException +from starlette.middleware import Middleware from starlette.requests import Request from starlette.responses import JSONResponse from starlette.routing import Route +from starlette.types import ASGIApp, Receive, Scope, Send from ._triggers import TriggerEngine, TriggerError from .extensions import ( @@ -63,10 +65,27 @@ from .tasks.v1 import ( TextPart as TaskTextPart, ) -from ..transfers import TransferError, upload +from ..transfers import TransferError, serving_on, upload logger = logging.getLogger(__name__) + +class _ServingOnItsPort: + """Each request is served ``serving_on`` the port this server took it on, so the transfers it makes + reach the server's own staging over loopback there, whatever ``A2A_PORT`` says.""" + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + server = scope.get("server") if scope["type"] == "http" else None + port = server[1] if server else None + if not isinstance(port, int): + await self.app(scope, receive, send) + return + with serving_on(port): + await self.app(scope, receive, send) + _AGENT_DEFINITION = "_agentenv_a2a_definition" _AgentT = TypeVar("_AgentT", bound="AgentEnvAgent") _MAX_CACHED_CONTEXTS = 1_024 @@ -918,7 +937,7 @@ def __init__(self, agent: Any, definition: _AgentDefinition | None = None) -> No ) if self.staging.max_bytes: routes.extend(staging_routes(self.staging)) - kwargs = {"routes": routes} + kwargs = {"routes": routes, "middleware": [Middleware(_ServingOnItsPort)]} if definition.lifespan is not None: kwargs["lifespan"] = definition.lifespan self.app = Starlette(**kwargs) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/transfers.py b/packages/agentenv-protocol/src/agentenv_protocol/transfers.py index 6deca71d..ef4fbd3b 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/transfers.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/transfers.py @@ -53,9 +53,10 @@ _T = TypeVar("_T") # A grant naming a path on its holder's own staging carries this header, whose value is that path on # the holder's server: a sandbox can't always call its own public URL, so the helpers reach the server -# over loopback, on the port agent-env deploys it on. +# over loopback. STAGING_PATH_HEADER = "AgentEnv-Staging-Path" _SERVER_PORT_ENV = "A2A_PORT" +_SERVER_PORT: ContextVar[int | None] = ContextVar("agentenv_protocol_server_port", default=None) def _https_url(value: str) -> str: @@ -286,12 +287,24 @@ def _is_s3_request_timeout(response: httpx.Response) -> bool: return False +@contextmanager +def serving_on(port: int) -> Iterator[None]: + """Transfers made in this scope serve a request the holder's own server took on ``port``, so they + reach its staging over loopback there. The SDK's server sets it for every request it handles.""" + + token = _SERVER_PORT.set(port) + try: + yield + finally: + _SERVER_PORT.reset(token) + + def loopback_url(url: str, headers: Mapping[str, str] | None) -> str | None: """Where the holder's own server answers a grant's ``url`` over loopback: when the grant names a path - on the holder's staging (an ``AgentEnv-Staging-Path`` header ending the URL's path) and ``A2A_PORT`` - says which port that server listens on. None otherwise.""" + on the holder's staging (an ``AgentEnv-Staging-Path`` header ending the URL's path) and the server's + port is known, from the request being served (``serving_on``) or else ``A2A_PORT``. None otherwise.""" path = httpx.Headers(headers or {}).get(STAGING_PATH_HEADER) - port = os.environ.get(_SERVER_PORT_ENV, "") + port = str(_SERVER_PORT.get() or os.environ.get(_SERVER_PORT_ENV, "")) if not path or not path.startswith("/") or not port.isdigit() or not 0 < int(port) < 65536: return None parsed = urlsplit(url) diff --git a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py index e7b1d36a..d2bcdb99 100644 --- a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py +++ b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py @@ -31,6 +31,7 @@ TaskObjectTrajectoryRequest, TaskRequest, TaskResult, + TaskTrajectoryRequest, TrajectoryWriteObjects, a2a_agent, card_request_accepts, @@ -1108,3 +1109,36 @@ def test_the_loopback_url_keeps_the_grants_query_and_reads_the_header_in_any_cas transfers.loopback_url(f"{_PUBLIC}{_STAGED}?part=1", headers) == f"http://127.0.0.1:8123{_STAGED}?part=1" ) + + +@a2a_agent(identity=AgentIdentity(name="own-port", description="test", version="1")) +class _UploadsItsTrajectory(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + return TaskResult.text("ok") + + @extension(TRAJECTORY_V1.get) + async def trajectory(self, request: TaskTrajectoryRequest | TaskObjectTrajectoryRequest): + uploaded = await upload(request.objects.trajectory, b"[]") + return {"objects": {"trajectory": uploaded.model_dump(exclude_none=True)}} + + +def test_the_sdk_server_reaches_its_staging_on_the_port_it_took_the_request_on( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("A2A_PORT", "9999") # not where this server listens + sent: list[str] = [] + + def handle(request: httpx.Request) -> httpx.Response: + sent.append(str(request.url)) + request.read() + return httpx.Response(201) + + _route(monkeypatch, handle) + write, _, _ = _on_own_staging(b"") + payload = {"task_id": "t", "objects": {"trajectory": write.model_dump(mode="json")}} + + with TestClient(_UploadsItsTrajectory().create_app(), base_url="http://127.0.0.1:8456") as client: + assert client.post("/ext/trajectory", json=payload).status_code == 200 + + assert sent == [f"http://127.0.0.1:8456{_STAGED}"] + assert transfers.loopback_url(write.write.url, write.write.headers) == f"http://127.0.0.1:9999{_STAGED}" diff --git a/src/agent_env/a2a_agent/object_transfer.py b/src/agent_env/a2a_agent/object_transfer.py index 5f1d3156..5f8826a6 100644 --- a/src/agent_env/a2a_agent/object_transfer.py +++ b/src/agent_env/a2a_agent/object_transfer.py @@ -61,8 +61,8 @@ _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. An agent on a current SDK reaches its staging over " - "loopback, on its A2A_PORT; an older one goes through its own public URL, which its sandbox provider may " - "not allow. An object store whose grants reach the agent avoids staging." + "loopback; an older one goes through its own public URL, which its sandbox provider may not allow. An " + "object store whose grants reach the agent avoids staging." ) # The time budget of one transfer, outermost first: a grant (the issuing store's diff --git a/tst/unit/a2a_agent/staging_test.py b/tst/unit/a2a_agent/staging_test.py index 8743f050..02c0dfec 100644 --- a/tst/unit/a2a_agent/staging_test.py +++ b/tst/unit/a2a_agent/staging_test.py @@ -59,7 +59,7 @@ from agent_env.task_step.context import TaskStepContext from tst.util.granting_object_store import GrantingObjectStore -URL = "https://agent.example.test" +URL = "https://agent.example.test:8443" # a port, so the agent's server knows the one it took a request on @a2a_agent(identity=AgentIdentity(name="staging-test", description="test", version="1")) @@ -141,7 +141,9 @@ def _serve(app, monkeypatch) -> None: @pytest.fixture(params=["public URL", "loopback"]) def agent(request, tmp_path, monkeypatch): """The agent's app, which every HTTP client in this process reaches at ``URL``: reaching its staging - through its public URL, or over loopback from a sandbox that can't call that URL.""" + through its public URL, as an agent whose helpers don't use loopback does, or over loopback from a + sandbox that can't call that URL. A transfer made while the app serves a request uses the port it took + the request on; the uploads these tests make themselves, as a hook process would, use ``A2A_PORT``.""" monkeypatch.setenv("AGENTENV_STAGING_DIR", str(tmp_path / "agent-staging")) _Agent.workdir = tmp_path / "agent" _Agent.received = {} @@ -151,7 +153,7 @@ def agent(request, tmp_path, monkeypatch): monkeypatch.setenv("A2A_PORT", "8000") app = _calling_its_own_url_fails(app) else: - monkeypatch.delenv("A2A_PORT", raising=False) + monkeypatch.setattr(transfers, "loopback_url", lambda url, headers: None) _serve(app, monkeypatch) return app @@ -349,9 +351,9 @@ def test_each_staged_grant_names_its_path_on_the_agents_own_server(store, monkey @pytest.mark.asyncio -async def test_an_agent_that_cant_call_its_own_url_and_has_no_loopback_port_fails_saying_why(tmp_path, monkeypatch, store): +async def test_an_agent_that_cant_call_its_own_url_and_doesnt_use_loopback_fails_saying_why(tmp_path, monkeypatch, store): monkeypatch.setenv("AGENTENV_STAGING_DIR", str(tmp_path / "agent-staging")) - monkeypatch.delenv("A2A_PORT", raising=False) + monkeypatch.setattr(transfers, "loopback_url", lambda url, headers: None) # an older SDK monkeypatch.setattr(transfers, "_RETRY_BACKOFF_SECONDS", 0) app = _calling_its_own_url_fails(_Agent().create_app()) _serve(app, monkeypatch) @@ -360,7 +362,7 @@ async def test_an_agent_that_cant_call_its_own_url_and_has_no_loopback_port_fail _, path = _method(card, TRAJECTORY_V1.uri, "get") target = store.object_url("trajectories/t.json") - with pytest.raises(httpx.HTTPStatusError, match="reaches its staging over loopback, on its A2A_PORT"): + with pytest.raises(httpx.HTTPStatusError, match="reaches its staging over loopback; an older one"): await fetch_trajectory(URL + path, {"task_id": "t"}, upload=TrajectoryUpload.to(granting, target), store=granting) assert store.get_object_metadata_at(target) is None