diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index fb719636..cd84562d 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -622,8 +622,18 @@ 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:{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 65dda9cc..ef4fbd3b 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,12 @@ _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. +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: @@ -281,6 +287,63 @@ 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 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 = 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) + 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 +464,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 +498,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 +556,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..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, @@ -964,3 +965,180 @@ 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" + ) + + +@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 4dcc2b37..5f8826a6 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; 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..02c0dfec 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 @@ -54,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")) @@ -109,19 +114,47 @@ 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, 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 = {} + _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.setattr(transfers, "loopback_url", lambda url, headers: None) + _serve(app, monkeypatch) return app @@ -302,6 +335,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_doesnt_use_loopback_fails_saying_why(tmp_path, monkeypatch, store): + monkeypatch.setenv("AGENTENV_STAGING_DIR", str(tmp_path / "agent-staging")) + 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) + 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; 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 + + @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")