fix(a2a): agents reach their own staging over loopback

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 |
 | --- | --- |

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)

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: