agentenv-framework-protocol 0.1.300__tar.gz → 0.1.302__tar.gz

{agentenv_framework_protocol-0.1.300 → agentenv_framework_protocol-0.1.302}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: agentenv-framework-protocol
-Version: 0.1.300
+Version: 0.1.302
 Summary: Open data-plane protocol and server SDK for agent environments.
 Project-URL: Homepage, https://www.agentenvframework.com
 Project-URL: Documentation, https://www.agentenvframework.com/docs
@@ -666,8 +666,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 |
 | --- | --- |

{agentenv_framework_protocol-0.1.300 → agentenv_framework_protocol-0.1.302}/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 |
 | --- | --- |

{agentenv_framework_protocol-0.1.300 → agentenv_framework_protocol-0.1.302}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "agentenv-framework-protocol"
-version = "0.1.300"
+version = "0.1.302"
 description = "Open data-plane protocol and server SDK for agent environments."
 readme = "README.md"
 license = "Apache-2.0"

{agentenv_framework_protocol-0.1.300 → agentenv_framework_protocol-0.1.302}/src/agentenv_protocol/a2a_agent/framework.py

@@ -22,9 +22,11 @@ from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
 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 (
 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 @@ class A2AAgentApplication:
                 )
         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)

{agentenv_framework_protocol-0.1.300 → agentenv_framework_protocol-0.1.302}/src/agentenv_protocol/transfers.py

@@ -10,14 +10,14 @@ import random
 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 @@ TRANSFER_STALL_BUDGET_SECONDS = TRANSFER_ATTEMPTS * (
     _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:

{agentenv_framework_protocol-0.1.300 → agentenv_framework_protocol-0.1.302}/tests/test_a2a_object_transfers.py

@@ -31,6 +31,7 @@ from agentenv_protocol.a2a_agent import (
     TaskObjectTrajectoryRequest,
     TaskRequest,
     TaskResult,
+    TaskTrajectoryRequest,
     TrajectoryWriteObjects,
     a2a_agent,
     card_request_accepts,
@@ -964,3 +965,180 @@ async def test_namespace_upload_of_bytes_declares_its_length(
         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}"