diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index ac57fcbb..05821584 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -583,3 +583,33 @@ SDK. A snapshot or changelog is restored in the form it was captured in: one captured as objects only through the object variants, an older one only through `s3_prefix`. Do not roll agent-env back once portable snapshots or changelogs exist, because older releases cannot load them. + +#### Staging + +When the object store's grants cannot reach the agent, as with a local store +and an agent on a remote sandbox, AgentEnv moves the objects through the agent +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. + +| Route | Does | +| --- | --- | +| `PUT /ext/staging/{path}` | stores the body | +| `GET /ext/staging/{path}` | returns it, with an `ETag` | +| `POST /ext/staging/{prefix}` | stores a multipart upload's `file` at `{prefix}/{key}`, `key` being a form field sent before it | +| `GET /ext/staging/{prefix}/` | lists `{objects: [{path, size_bytes, etag}]}` below the prefix | +| `DELETE /ext/staging/{path}` | removes the object, only while it still matches an `If-Match` tag when one is given | +| `DELETE /ext/staging/{prefix}/` | removes everything below the prefix | + +A path's first segment is an id AgentEnv generates, at least 22 characters, and +never shares, which is what keeps staged objects private on a public agent URL. +Everything staged counts against `AGENTENV_STAGING_MAX_BYTES` (16 GiB unless +set; `0` turns staging off) and lives until AgentEnv removes it or the server +process ends, in a directory of the process's own (or under +`AGENTENV_STAGING_DIR` when set) that only the server's user can read. One +server process owns a staging directory. An agent that doesn't use the SDK can serve the same routes +and add `{uri: "urn:agentenv:staging/v1", params: {endpoint: "/ext/staging"}}` +to its card. diff --git a/packages/agentenv-protocol/THIRD_PARTY_NOTICES.md b/packages/agentenv-protocol/THIRD_PARTY_NOTICES.md index 7bbbd90c..4126ccac 100644 --- a/packages/agentenv-protocol/THIRD_PARTY_NOTICES.md +++ b/packages/agentenv-protocol/THIRD_PARTY_NOTICES.md @@ -191,6 +191,13 @@ The runtime dependencies of `agentenv-protocol` and its `agent` extra, at the ve - Source: - [License text 2](#license-text-2) +### python-multipart 0.0.32 + +- License: Apache-2.0 +- Author: Andrew Dunham +- Source: +- [License text 1](#license-text-1) + ### regex 2026.6.28 - License: Apache-2.0 AND CNRI-Python diff --git a/packages/agentenv-protocol/pyproject.toml b/packages/agentenv-protocol/pyproject.toml index 3f6775a2..60df9ec5 100644 --- a/packages/agentenv-protocol/pyproject.toml +++ b/packages/agentenv-protocol/pyproject.toml @@ -16,11 +16,13 @@ dependencies = [ [project.optional-dependencies] agent = [ "a2a-sdk[http-server]>=0.3.26,<0.4", + "python-multipart>=0.0.13", # the staging extension's upload forms "regex>=2023.12.25", "uvicorn>=0.30", ] dev = [ "a2a-sdk[http-server]>=0.3.26,<0.4", + "python-multipart>=0.0.13", "regex>=2023.12.25", "pytest>=9.0", "pytest-asyncio>=0.23.0", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py index f62bbf5b..1e978a33 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py @@ -76,6 +76,7 @@ RegisteredOperation, build_registry, ) +from .staging import STAGING_ENDPOINT, STAGING_V1_URI, StagingStore, staging_routes from .tasks.v1 import ( AgentConfig, AgentRunResult, @@ -116,6 +117,8 @@ "PEER_AGENTS_V1", "SKILL_CONFIG_V1", "SNAPSHOT_V1", + "STAGING_ENDPOINT", + "STAGING_V1_URI", "STANDARD_EXTENSIONS", "TRAJECTORY_V1", "TRIGGERS_V1", @@ -172,6 +175,7 @@ "SnapshotReadObjects", "SnapshotUploadedObjects", "SnapshotWriteObjects", + "StagingStore", "TaskError", "TaskObjectTrajectoryRequest", "TaskOutcome", @@ -204,5 +208,6 @@ "extension", "request_fields", "serve", + "staging_routes", "upload", ] 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 54495f33..0ca4b06b 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py @@ -45,6 +45,7 @@ enable, ) from .registry import RegisteredOperation, build_registry +from .staging import STAGING_ENDPOINT, StagingStore, staging_routes from .tasks.v1 import ( AgentConfig, AgentRunResult, @@ -873,6 +874,15 @@ def __init__(self, agent: Any, definition: _AgentDefinition | None = None) -> No ("/.well-known/agent-card.json", "GET"): "Agent Card", } ) + # Objects agent-env moves through the agent when its object store's grants cannot reach it. + self.staging = StagingStore() + for extension in self.registry.extensions if self.staging.max_bytes else (): + for operation in extension.operations.values(): + if operation.definition.path.startswith(STAGING_ENDPOINT + "/"): + raise ValueError( + f"extension operation {extension.definition.uri}.{operation.definition.name} " + f"conflicts with the staging routes below {STAGING_ENDPOINT}" + ) conformance = self.registry.conformance() for override in conformance["standard_operation_overrides"]: logger.warning( @@ -906,6 +916,8 @@ def __init__(self, agent: Any, definition: _AgentDefinition | None = None) -> No methods=[operation.definition.method], ) ) + if self.staging.max_bytes: + routes.extend(staging_routes(self.staging)) kwargs = {"routes": routes} if definition.lifespan is not None: kwargs["lifespan"] = definition.lifespan @@ -953,6 +965,8 @@ def _build_card(self) -> Any: AgentExtension.model_validate(item) for item in self.registry.card_extensions() ] + if self.staging.max_bytes: + extensions.append(AgentExtension.model_validate(self.staging.card_extension())) resolved_capabilities = AgentCapabilities( streaming=self.streaming, push_notifications=False, diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/staging.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/staging.py new file mode 100644 index 00000000..46018e9a --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/staging.py @@ -0,0 +1,512 @@ +"""The staging extension: a small object store on the agent's own server. + +When an object store's grants cannot reach an agent's sandbox, as with a local store and a remote +sandbox, agent-env moves the objects through the agent instead. Before a call it pushes what the agent +will read into staging; after the call it pulls what the agent wrote. The grants it sends are ordinary +HTTPS grants naming staged paths on the agent's own URL, so the agent's transfer helpers don't change. + +Routes, below the extension's endpoint: + +* ``PUT {path}`` stores the request body; +* ``GET {path}`` returns it, with an ``ETag``; +* ``POST {prefix}`` stores a multipart upload's file at ``{prefix}/{key}``, the way an upload-policy + grant is used: form fields first, ``key`` among them, then ``file``; +* ``GET {prefix}/`` lists what is stored below a prefix; +* ``DELETE {path}`` removes an object, and only while it still has the ``If-Match`` tag when one is + given; ``DELETE {prefix}/`` removes everything below a prefix. + +A path's first segment is an id agent-env generates and tells no one else, so it must be long enough +not to be guessed, and the directories holding staged objects are the server user's alone. Everything +staged counts against ``AGENTENV_STAGING_MAX_BYTES`` (a limit of 0 turns staging off), and lives until +agent-env removes it or the server process ends: in a directory of the process's own, or under +``AGENTENV_STAGING_DIR`` when that is set. One server process owns a staging directory, since the byte +count and the ``If-Match`` check are its own. +""" + +from __future__ import annotations + +import asyncio +import atexit +import hashlib +import os +import shutil +import tempfile +import time +from collections.abc import AsyncIterator +from dataclasses import dataclass +from pathlib import Path +from stat import S_ISREG +from typing import BinaryIO + +from python_multipart.exceptions import FormParserError +from python_multipart.multipart import MultipartParser, parse_options_header +from starlette.requests import ClientDisconnect, Request +from starlette.responses import JSONResponse, Response, StreamingResponse +from starlette.routing import Route + +STAGING_V1_URI = "urn:agentenv:staging/v1" +STAGING_ENDPOINT = "/ext/staging" +STAGING_DIR_ENV = "AGENTENV_STAGING_DIR" +STAGING_MAX_BYTES_ENV = "AGENTENV_STAGING_MAX_BYTES" +DEFAULT_STAGING_MAX_BYTES = 16 * 1024**3 +MIN_ID_LENGTH = 22 # about 128 bits of a URL-safe base64 id + +_CHUNK_BYTES = 1024 * 1024 +_MAX_PATH_LENGTH = 1024 +_KEY_FIELD = "key" +_FILE_FIELD = "file" +_MAX_PART_HEADER_BYTES = 16 * 1024 +_MAX_FIELD_BYTES = 64 * 1024 +_MAX_FIELDS = 32 + + +class StagingError(Exception): + def __init__(self, status_code: int, message: str) -> None: + super().__init__(message) + self.status_code = status_code + + +@dataclass(frozen=True) +class Staged: + """One stored object, as a write or a listing reports it.""" + + path: str + size_bytes: int + etag: str + sha256: str | None = None + + def to_json(self) -> dict[str, object]: + body: dict[str, object] = {"path": self.path, "size_bytes": self.size_bytes, "etag": self.etag} + if self.sha256 is not None: + body["sha256"] = self.sha256 + return body + + +def staging_max_bytes() -> int: + """The configured staging limit, ``AGENTENV_STAGING_MAX_BYTES`` or 16 GiB.""" + value = os.environ.get(STAGING_MAX_BYTES_ENV, "").strip() + if not value: + return DEFAULT_STAGING_MAX_BYTES + limit = int(value) + if limit < 0: + raise ValueError(f"{STAGING_MAX_BYTES_ENV} must not be negative") + return limit + + +class StagingStore: + """The objects staged on one agent, below ``root``, within ``max_bytes`` in all. + + Writes land in a staging file and are renamed into place, so a reader sees a whole object or none. + Each write stamps a unique modification time, which with the size makes the object's tag. + """ + + def __init__(self, root: str | Path | None = None, *, max_bytes: int | None = None) -> None: + configured = root or os.environ.get(STAGING_DIR_ENV) + self._root = Path(configured) if configured else None + self.max_bytes = staging_max_bytes() if max_bytes is None else max_bytes + self._lock = asyncio.Lock() + self._used: int | None = None # stored plus reserved bytes, counted on first use + self._last_stamp = 0 + + @property + def root(self) -> Path: + """Where staged objects live; unless one was given, a new directory, made on first use and removed + when the process exits, so a restarted server never counts what an earlier one left.""" + if self._root is None: + self._root = Path(tempfile.mkdtemp(prefix="agentenv-staging-")) + atexit.register(shutil.rmtree, self._root, ignore_errors=True) + return self._root + + @property + def _objects(self) -> Path: + return self.root / "objects" + + @property + def _incoming(self) -> Path: + return self.root / "incoming" + + def card_extension(self, endpoint: str = STAGING_ENDPOINT) -> dict[str, object]: + return { + "uri": STAGING_V1_URI, + "description": "Stage objects on the agent for agent-env to push and pull.", + "params": {"endpoint": endpoint, "max_bytes": self.max_bytes}, + } + + async def put(self, path: str, chunks: AsyncIterator[bytes], *, size_bytes: int | None = None) -> Staged: + """Store ``chunks`` at ``path``, replacing what is there. ``size_bytes``, when known, is + reserved before the first byte arrives and must be what arrives.""" + target = self._file(path) + async with _Incoming(self, size_bytes) as incoming: + async for chunk in chunks: + await incoming.write(chunk) + if size_bytes is not None and incoming.size_bytes != size_bytes: + raise StagingError(400, "The body is shorter than its Content-Length.") + return await self._commit(incoming, target, path) + + def open(self, path: str) -> tuple[BinaryIO, Staged]: + """The object at ``path``, opened, and its description.""" + try: + file = self._file(path).open("rb") + except (FileNotFoundError, IsADirectoryError, NotADirectoryError): + raise StagingError(404, "Nothing is staged at this path.") from None + stat = os.fstat(file.fileno()) + return file, Staged(path, stat.st_size, _etag(stat)) + + def list_below(self, prefix: str) -> list[Staged]: + """Everything below ``prefix``, by path relative to it.""" + top = self._file(prefix) + found = [] + for directory, _, files in os.walk(top): + for name in files: + file = Path(directory) / name + try: + stat = file.stat() + except FileNotFoundError: + continue + found.append(Staged(file.relative_to(top).as_posix(), stat.st_size, _etag(stat))) + return sorted(found, key=lambda staged: staged.path) + + async def delete(self, path: str, *, if_match: str | None = None) -> None: + target = self._file(path) + async with self._lock: + used = await self._usage() + try: + found = target.stat() + except (FileNotFoundError, NotADirectoryError): + raise StagingError(404, "Nothing is staged at this path.") from None + if not S_ISREG(found.st_mode): + raise StagingError(404, "Nothing is staged at this path.") + if if_match is not None and if_match != _etag(found): + raise StagingError(412, "The staged object has changed.") + target.unlink() + self._used = used - found.st_size + self._prune(target.parent) + + async def delete_below(self, prefix: str) -> None: + top = self._file(prefix) + async with self._lock: + used = await self._usage() + freed = sum(staged.size_bytes for staged in await asyncio.to_thread(self.list_below, prefix)) + await asyncio.to_thread(shutil.rmtree, top, ignore_errors=True) + self._used = used - freed + self._prune(top.parent) + + async def post(self, prefix: str, request: Request) -> Staged: + """Store a multipart upload's file at ``{prefix}/{key}``.""" + prefix_path = _check_path(prefix) + media_type, params = parse_options_header(request.headers.get("content-type")) + if media_type != b"multipart/form-data" or b"boundary" not in params: + raise StagingError(400, "The upload must be multipart/form-data.") + form = _Form() + parser = MultipartParser(params[b"boundary"], form.callbacks()) + async with _Incoming(self, None) as incoming: + try: + async for chunk in request.stream(): + parser.write(chunk) + await incoming.write(form.take_file_bytes()) + parser.finalize() + except FormParserError: + raise StagingError(400, "The upload is not a valid multipart form.") from None + await incoming.write(form.take_file_bytes()) + if not form.file_done: + raise StagingError(400, f"The upload has no {_FILE_FIELD!r} field.") + path = _check_path(f"{prefix_path}/{form.key}") + return await self._commit(incoming, self._file(path), path) + + def _file(self, path: str) -> Path: + return self._objects / _check_path(path) + + async def _reserve(self, size_bytes: int) -> None: + async with self._lock: + used = await self._usage() + if used + size_bytes > self.max_bytes: + raise StagingError(413, "Staging is full.") + self._used = used + size_bytes + + async def _release(self, size_bytes: int) -> None: + async with self._lock: + self._used = await self._usage() - size_bytes + + async def _usage(self) -> int: + if self._used is None: + self._used = await asyncio.to_thread(_tree_bytes, self._objects) + return self._used + + async def _commit(self, incoming: _Incoming, target: Path, path: str) -> Staged: + await incoming.close() + async with self._lock: + used = await self._usage() + try: + replaced = target.stat().st_size if target.is_file() else 0 + _private(self._objects) + target.parent.mkdir(parents=True, exist_ok=True) + self._stamp(incoming.path) + os.replace(incoming.path, target) + except (FileExistsError, IsADirectoryError, NotADirectoryError): + raise StagingError(409, "The path conflicts with one already staged.") from None + incoming.committed() + self._used = used - replaced + stat = target.stat() + return Staged(path, stat.st_size, _etag(stat), incoming.sha256) + + def _stamp(self, file: Path) -> None: + stamp = max(time.time_ns(), self._last_stamp + 1000) + self._last_stamp = stamp + os.utime(file, ns=(stamp, stamp)) + + def _prune(self, directory: Path) -> None: + while directory != self._objects and self._objects in directory.parents: + try: + directory.rmdir() + except OSError: + return + directory = directory.parent + + +class _Incoming: + """A write on its way in: a staging file, the bytes it reserves, and their digest.""" + + def __init__(self, store: StagingStore, size_bytes: int | None) -> None: + self._store = store + self._declared = size_bytes + self._reserved = 0 + self._file: BinaryIO | None = None + self._committed = False + self.path = Path() + self.size_bytes = 0 + self._digest = hashlib.sha256() + self._buffer = bytearray() + + async def __aenter__(self) -> _Incoming: + if self._declared is not None: + await self._store._reserve(self._declared) + self._reserved = self._declared + try: + await asyncio.to_thread(_private, self._store._incoming) + fd, name = tempfile.mkstemp(dir=self._store._incoming) + except BaseException: + await self._store._release(self._reserved) + raise + self._file = os.fdopen(fd, "wb") + self.path = Path(name) + return self + + async def __aexit__(self, *exc_info: object) -> None: + if self._file is not None: + self._file.close() + if not self._committed: + self.path.unlink(missing_ok=True) + await self._store._release(self._reserved) + + @property + def sha256(self) -> str: + return self._digest.hexdigest() + + async def write(self, data: bytes) -> None: + if not data: + return + self.size_bytes += len(data) + if self._declared is not None and self.size_bytes > self._declared: + raise StagingError(400, "The body is longer than its Content-Length.") + if self.size_bytes > self._reserved: + await self._store._reserve(self.size_bytes - self._reserved) + self._reserved = self.size_bytes + self._digest.update(data) + self._buffer += data + if len(self._buffer) >= _CHUNK_BYTES: + await self._flush() + + async def _flush(self) -> None: + data, self._buffer = bytes(self._buffer), bytearray() + await asyncio.to_thread(self._file.write, data) + + async def close(self) -> None: + await self._flush() + await asyncio.to_thread(self._file.close) + self._file = None + + def committed(self) -> None: + """The staging file is now the object: its bytes stay counted, and it stays.""" + self._committed = True + + +class _Form: + """A multipart form read as an upload policy is: fields, then the file, whose bytes are handed + out as they arrive. Parts after the file are ignored.""" + + def __init__(self) -> None: + self._fields: dict[str, bytes] = {} + self._headers: dict[bytes, bytes] = {} + self._header_bytes = 0 + self._header_name = b"" + self._header_value = b"" + self._part_name: str | None = None + self._part_data = bytearray() + self._file_bytes = bytearray() + self.key = "" + self.file_done = False + + def callbacks(self) -> dict: + return { + "on_part_begin": self._on_part_begin, + "on_header_field": self._on_header_field, + "on_header_value": self._on_header_value, + "on_header_end": self._on_header_end, + "on_headers_finished": self._on_headers_finished, + "on_part_data": self._on_part_data, + "on_part_end": self._on_part_end, + } + + def take_file_bytes(self) -> bytes: + data, self._file_bytes = bytes(self._file_bytes), bytearray() + return data + + def _on_part_begin(self) -> None: + self._headers = {} + self._header_bytes = 0 + self._part_name = None + self._part_data = bytearray() + + def _on_header_field(self, data: bytes, start: int, end: int) -> None: + self._count_header_bytes(end - start) + self._header_name += data[start:end] + + def _on_header_value(self, data: bytes, start: int, end: int) -> None: + self._count_header_bytes(end - start) + self._header_value += data[start:end] + + def _count_header_bytes(self, count: int) -> None: + self._header_bytes += count + if self._header_bytes > _MAX_PART_HEADER_BYTES: + raise StagingError(400, "A form part's headers are too large.") + + def _on_header_end(self) -> None: + self._headers[self._header_name.lower()] = self._header_value + self._header_name = self._header_value = b"" + + def _on_headers_finished(self) -> None: + if self.file_done: + return + _, options = parse_options_header(self._headers.get(b"content-disposition")) + name = options.get(b"name") + if name is None: + raise StagingError(400, "A form field has no name.") + self._part_name = name.decode("utf-8", "replace") + if self._part_name == _FILE_FIELD: + key = self._fields.get(_KEY_FIELD) + if key is None: + raise StagingError(400, f"The upload must name its {_KEY_FIELD!r} before its file.") + self.key = key.decode("utf-8", "replace") + elif len(self._fields) >= _MAX_FIELDS: + raise StagingError(400, "The upload has too many form fields.") + + def _on_part_data(self, data: bytes, start: int, end: int) -> None: + if self.file_done: + return + if self._part_name == _FILE_FIELD: + self._file_bytes += data[start:end] + return + if len(self._part_data) + end - start > _MAX_FIELD_BYTES: + raise StagingError(400, "A form field is too large.") + self._part_data += data[start:end] + + def _on_part_end(self) -> None: + if self.file_done: + return + if self._part_name == _FILE_FIELD: + self.file_done = True + elif self._part_name is not None: + self._fields[self._part_name] = bytes(self._part_data) + + +def _check_path(path: str) -> str: + """``path`` when it is a normalized relative path whose first segment is a long enough id.""" + parts = path.split("/") + if ( + len(path) > _MAX_PATH_LENGTH + or "\\" in path + or "\x00" in path + or any(part in ("", ".", "..") for part in parts) + ): + raise StagingError(400, "The staging path is not a normalized relative path.") + if len(parts[0]) < MIN_ID_LENGTH: + raise StagingError(404, "Nothing is staged at this path.") + return path + + +def _private(directory: Path) -> None: + """Make ``directory`` readable by its owner alone, so other users on the host cannot list the ids + below it, however it was left. Only staging's own directories, never the one it was given.""" + directory.mkdir(parents=True, exist_ok=True) + if directory.stat().st_mode & 0o077: + directory.chmod(0o700) + + +def _etag(stat: os.stat_result) -> str: + return f'"{stat.st_mtime_ns:x}-{stat.st_size:x}"' + + +def _tree_bytes(top: Path) -> int: + total = 0 + for directory, _, files in os.walk(top): + for name in files: + try: + total += (Path(directory) / name).stat().st_size + except FileNotFoundError: + pass + return total + + +def staging_routes(store: StagingStore, endpoint: str = STAGING_ENDPOINT) -> list[Route]: + """The staging extension's routes, below ``endpoint``.""" + + async def handle(request: Request) -> Response: + path: str = request.path_params["path"] + listing = path.endswith("/") + try: + if listing: + path = path[:-1] + if request.method == "GET": + found = await asyncio.to_thread(store.list_below, path) + return JSONResponse({"objects": [staged.to_json() for staged in found]}) + if request.method == "DELETE": + await store.delete_below(path) + return Response(status_code=204) + raise StagingError(405, "A prefix can only be listed or deleted.") + if request.method == "GET": + file, staged = await asyncio.to_thread(store.open, path) + return StreamingResponse( + _read(file), + media_type="application/octet-stream", + headers={"content-length": str(staged.size_bytes), "etag": staged.etag}, + ) + if request.method == "PUT": + length = request.headers.get("content-length") + staged = await store.put( + path, request.stream(), size_bytes=int(length) if length is not None else None + ) + return JSONResponse(staged.to_json(), status_code=201) + if request.method == "POST": + staged = await store.post(path, request) + return JSONResponse(staged.to_json(), status_code=201) + await store.delete(path, if_match=request.headers.get("if-match")) + return Response(status_code=204) + except StagingError as exc: + return JSONResponse({"detail": str(exc)}, status_code=exc.status_code) + except ValueError: + return JSONResponse({"detail": "The request is malformed."}, status_code=400) + except ClientDisconnect: + return Response(status_code=400) + except OSError as exc: + status = 507 if exc.errno == 28 else 500 # ENOSPC: the sandbox's disk is full + return JSONResponse({"detail": "Staging could not store the object."}, status_code=status) + + return [Route(f"{endpoint}/{{path:path}}", handle, methods=["GET", "PUT", "POST", "DELETE"])] + + +async def _read(file: BinaryIO) -> AsyncIterator[bytes]: + try: + while chunk := await asyncio.to_thread(file.read, _CHUNK_BYTES): + yield chunk + finally: + file.close() diff --git a/packages/agentenv-protocol/src/agentenv_protocol/transfers.py b/packages/agentenv-protocol/src/agentenv_protocol/transfers.py index 42521dd0..65dda9cc 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/transfers.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/transfers.py @@ -225,7 +225,7 @@ def filter(self, record: logging.LogRecord) -> bool: @contextmanager -def _redacting_request_urls() -> Iterator[None]: +def redacting_request_urls() -> Iterator[None]: """httpx logs only the origin of requests made in this scope: grant URLs carry credentials.""" token = _REDACT_REQUEST_URLS.set(True) @@ -255,7 +255,7 @@ def _as_transfer_errors(expires_at: datetime) -> Iterator[None]: def _transfer_request(expires_at: datetime) -> Iterator[None]: """One request made with a grant.""" - with _redacting_request_urls(), _as_transfer_errors(expires_at): + with redacting_request_urls(), _as_transfer_errors(expires_at): yield diff --git a/packages/agentenv-protocol/tests/test_a2a_agent.py b/packages/agentenv-protocol/tests/test_a2a_agent.py index 56833d75..65cd6287 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent.py @@ -16,6 +16,7 @@ from a2a.types import ( AgentCapabilities as UpstreamAgentCapabilities, ) +from a2a.types import AgentExtension from a2a.types import ( Message, MessageSendParams, @@ -35,6 +36,7 @@ PEER_AGENTS_V1, SKILL_CONFIG_V1, SNAPSHOT_V1, + STAGING_V1_URI, TRAJECTORY_V1, TRIGGERS_V1, AgentCapabilities, @@ -62,6 +64,7 @@ PeerAgentsSetRequest, RequestDefinition, RequestVariant, + StagingStore, TaskObjectTrajectoryRequest, TaskOutcome, TaskProgress, @@ -694,7 +697,7 @@ async def run(self, request: TaskRequest) -> TaskResult: streaming=False, push_notifications=False, state_transition_history=False, - extensions=[], + extensions=[AgentExtension.model_validate(StagingStore().card_extension())], ) with pytest.raises(TypeError, match="unexpected keyword argument 'capabilities'"): @@ -731,7 +734,7 @@ async def run(self, request: TaskRequest): with TestClient(Agent().create_app()) as client: card = client.get("/.well-known/agent.json").json() assert card["capabilities"] == { - "extensions": [], + "extensions": [StagingStore().card_extension()], "pushNotifications": False, "stateTransitionHistory": False, "streaming": True, @@ -2625,6 +2628,7 @@ async def add_bundle_skill(self, request: BundleSkillRequest): "urn:agentenv:skill-config/v1", "urn:agentenv:trajectory/v1", "urn:agentenv:triggers/v1", + STAGING_V1_URI, } assert ( diff --git a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py index e0377e7b..ba10fe4e 100644 --- a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py +++ b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py @@ -847,7 +847,7 @@ def test_grant_urls_are_redacted_however_httpx_formats_its_log_call( ) -> None: caplog.set_level(logging.INFO, logger="httpx") - with transfers._redacting_request_urls(): + with transfers.redacting_request_urls(): logging.getLogger("httpx").info(message, *args) (record,) = caplog.records diff --git a/packages/agentenv-protocol/tests/test_a2a_staging.py b/packages/agentenv-protocol/tests/test_a2a_staging.py new file mode 100644 index 00000000..68ab1c6f --- /dev/null +++ b/packages/agentenv-protocol/tests/test_a2a_staging.py @@ -0,0 +1,275 @@ +from __future__ import annotations + +import hashlib +import os +import shutil +import subprocess +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import httpx +import pytest +from agentenv_protocol import transfers +from agentenv_protocol.a2a_agent import ( + STAGING_V1_URI, + AgentEnvAgent, + AgentIdentity, + StagingStore, + TaskRequest, + TaskResult, + a2a_agent, + custom_extension, + staging_routes, +) +from agentenv_protocol.transfers import ( + HttpGetGrant, + HttpPostPolicyGrant, + HttpPutGrant, + NamespaceUploader, + ReadObject, + WriteNamespaceGrant, + WriteObject, + download, + upload, +) +from starlette.applications import Starlette +from starlette.testclient import TestClient + +_UTC = timezone.utc # noqa: UP017 -- package supports Python 3.10. +SESSION = "s" * 32 +BASE = f"/ext/staging/{SESSION}" + + +@a2a_agent(identity=AgentIdentity(name="staging", description="test", version="1")) +class Agent(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + return TaskResult.text("ok") + + +@pytest.fixture(autouse=True) +def _staging_dir(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AGENTENV_STAGING_DIR", str(tmp_path / "staging")) + + +def _app(max_bytes: int = 1024) -> Starlette: + return Starlette(routes=staging_routes(StagingStore(max_bytes=max_bytes))) + + +def _expiry() -> datetime: + return datetime.now(_UTC) + timedelta(minutes=5) + + +def test_an_sdk_agent_advertises_staging_and_serves_it() -> None: + with TestClient(Agent().create_app()) as client: + card = client.get("/.well-known/agent-card.json").json() + stored = client.put(f"{BASE}/a/b.txt", content=b"hello") + fetched = client.get(f"{BASE}/a/b.txt") + listed = client.get(f"{BASE}/") + deleted = client.delete(f"{BASE}/a/b.txt") + gone = client.get(f"{BASE}/a/b.txt") + + [staging] = [e for e in card["capabilities"]["extensions"] if e["uri"] == STAGING_V1_URI] + assert staging["params"]["endpoint"] == "/ext/staging" + assert stored.status_code == 201 + assert stored.json()["sha256"] == hashlib.sha256(b"hello").hexdigest() + assert (fetched.content, fetched.headers["etag"]) == (b"hello", stored.json()["etag"]) + assert listed.json() == { + "objects": [{"path": "a/b.txt", "size_bytes": 5, "etag": stored.json()["etag"]}] + } + assert (deleted.status_code, gone.status_code) == (204, 404) + + +def test_a_limit_of_zero_turns_staging_off(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AGENTENV_STAGING_MAX_BYTES", "0") + with TestClient(Agent().create_app()) as client: + card = client.get("/.well-known/agent-card.json").json() + stored = client.put(f"{BASE}/a", content=b"x") + assert STAGING_V1_URI not in {e["uri"] for e in card["capabilities"]["extensions"]} + assert stored.status_code == 404 + + +def test_staging_off_leaves_its_paths_to_the_agent(monkeypatch: pytest.MonkeyPatch) -> None: + @a2a_agent(identity=AgentIdentity(name="own-route", description="test", version="1")) + class OwnRoute(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + return TaskResult.text("ok") + + @custom_extension(uri="urn:example:own/v1", operation="read", method="GET", path="/ext/staging/own") + async def read(self): + return {"status": "ok"} + + with pytest.raises(ValueError, match="conflicts with the staging routes"): + OwnRoute().create_app() + monkeypatch.setenv("AGENTENV_STAGING_MAX_BYTES", "0") + with TestClient(OwnRoute().create_app()) as client: + assert client.get("/ext/staging/own").json() == {"status": "ok"} + + +def test_staged_objects_are_the_server_users_alone(tmp_path) -> None: + with TestClient(_app()) as client: + client.put(f"{BASE}/a", content=b"x") + for directory in ("objects", "incoming"): + assert (tmp_path / "staging" / directory).stat().st_mode & 0o777 == 0o700 + + +def test_staging_tightens_its_own_directories_left_open(tmp_path) -> None: + for directory in ("objects", "incoming"): + (tmp_path / "staging" / directory).mkdir(parents=True) + (tmp_path / "staging" / directory).chmod(0o755) + with TestClient(_app()) as client: + client.put(f"{BASE}/a", content=b"x") + for directory in ("objects", "incoming"): + assert (tmp_path / "staging" / directory).stat().st_mode & 0o777 == 0o700 + + +def test_without_a_directory_each_server_stages_in_its_own(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AGENTENV_STAGING_DIR") + first, second = StagingStore(max_bytes=4), StagingStore(max_bytes=4) + with TestClient(Starlette(routes=staging_routes(first))) as client: + assert client.put(f"{BASE}/a", content=b"1234").status_code == 201 + with TestClient(Starlette(routes=staging_routes(second))) as client: # what the first left is not counted + assert client.put(f"{BASE}/a", content=b"1234").status_code == 201 + assert client.get(f"{BASE}/a").content == b"1234" + assert first.root != second.root and first.root.stat().st_mode & 0o777 == 0o700 + shutil.rmtree(first.root) + shutil.rmtree(second.root) + + +def test_a_servers_own_staging_directory_goes_when_it_exits(tmp_path) -> None: + script = ( + "from agentenv_protocol.a2a_agent import StagingStore\n" + "store = StagingStore(max_bytes=1)\n" + f"open({str(tmp_path / 'root')!r}, 'w').write(str(store.root))\n" + ) + env = {k: v for k, v in os.environ.items() if k != "AGENTENV_STAGING_DIR"} + subprocess.run([sys.executable, "-c", script], check=True, env=env) + assert not Path((tmp_path / "root").read_text()).exists() + + +def test_a_delete_spares_an_object_rewritten_since_it_was_read() -> None: + with TestClient(_app()) as client: + client.put(f"{BASE}/n", content=b"one") + read_tag = client.get(f"{BASE}/n").headers["etag"] + client.put(f"{BASE}/n", content=b"two") + stale = client.delete(f"{BASE}/n", headers={"if-match": read_tag}) + kept = client.get(f"{BASE}/n") + current = client.delete(f"{BASE}/n", headers={"if-match": kept.headers["etag"]}) + assert (stale.status_code, kept.content, current.status_code) == (412, b"two", 204) + + +@pytest.mark.parametrize( + "path, status", + [ + ("short/a", 404), # an id too short to be unguessable + (f"{SESSION}/%2E%2E/x", 400), # sent encoded, as clients fold a literal ".." + (f"{SESSION}//x", 400), + (f"{SESSION}/a\\b", 400), + ], +) +def test_staging_refuses_guessable_and_unnormalized_paths(path: str, status: int) -> None: + with TestClient(_app()) as client: + assert client.put(f"/ext/staging/{path}", content=b"x").status_code == status + + +def test_staging_lists_nothing_without_an_id() -> None: + with TestClient(_app()) as client: + client.put(f"{BASE}/a", content=b"x") + assert client.get("/ext/staging/").status_code == 400 + + +def test_staging_holds_to_its_limit_and_frees_what_it_drops() -> None: + with TestClient(_app(max_bytes=10)) as client: + assert client.put(f"{BASE}/a", content=b"123456").status_code == 201 + assert client.put(f"{BASE}/b", content=b"123456").status_code == 413 + # A replacement needs room beside the old copy while it lands, then frees the old one. + assert client.put(f"{BASE}/a", content=b"1234").status_code == 201 + assert client.put(f"{BASE}/b", content=b"123456").status_code == 201 + client.delete(f"{BASE}/a") + assert client.put(f"{BASE}/c", content=b"1234").status_code == 201 + client.delete(f"{BASE}/") + assert client.put(f"{BASE}/d", content=b"1234567890").status_code == 201 + + +def test_a_write_that_outgrows_the_limit_leaves_nothing_behind(tmp_path) -> None: + with TestClient(_app(max_bytes=4)) as client: + streamed = client.put(f"{BASE}/a", content=iter([b"12", b"345"])) # no Content-Length + assert (streamed.status_code, client.get(f"{BASE}/a").status_code) == (413, 404) + assert client.put(f"{BASE}/a", content=b"1234").status_code == 201 # its reservation was released + assert not any((tmp_path / "staging" / "incoming").iterdir()) + + +def test_a_path_cannot_be_both_an_object_and_a_prefix() -> None: + with TestClient(_app()) as client: + client.put(f"{BASE}/a", content=b"x") + assert client.put(f"{BASE}/a/b", content=b"x").status_code == 409 + + +def test_an_upload_must_name_its_key_before_its_file() -> None: + with TestClient(_app()) as client: + missing = client.post(BASE, files={"file": ("f", b"x")}) + escaping = client.post(BASE, data={"key": "../x"}, files={"file": ("f", b"x")}) + assert (missing.status_code, escaping.status_code) == (400, 400) + + +@pytest.mark.asyncio +async def test_the_transfer_helpers_move_objects_through_staging( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + app = _app(max_bytes=1 << 20) + real_async_client = httpx.AsyncClient + monkeypatch.setattr( + transfers.httpx, + "AsyncClient", + lambda **kwargs: real_async_client(transport=httpx.ASGITransport(app=app), **kwargs), + ) + monkeypatch.setattr( + transfers.httpx, + "Client", + lambda **kwargs: TestClient(app, follow_redirects=False), + ) + url = f"https://agent.example.test{BASE}" + body = b"x" * 300_000 + + uploaded = await upload( + WriteObject( + media_type="application/octet-stream", + max_bytes=len(body), + write=HttpPutGrant(kind="http-put", url=f"{url}/0", expires_at=_expiry()), + ), + body, + ) + await download( + ReadObject( + media_type="application/octet-stream", + max_bytes=len(body), + size_bytes=len(body), + sha256=uploaded.sha256, + read=HttpGetGrant(kind="http-get", url=f"{url}/0", expires_at=_expiry()), + ), + tmp_path / "out", + ) + namespace = NamespaceUploader( + WriteNamespaceGrant( + root_path="changelog/run-1", + expires_at=_expiry(), + max_objects=10, + max_object_bytes=1024, + max_total_bytes=4096, + write=HttpPostPolicyGrant( + kind="http-post-policy", + url=f"{url}/ns", + fields={}, + path_field="key", + file_field="file", + ), + ) + ) + await namespace.upload("000001.tar", b"increment") + + assert (tmp_path / "out").read_bytes() == body + with TestClient(app) as client: + listed = client.get(f"{BASE}/ns/").json()["objects"] + increment = client.get(f"{BASE}/ns/changelog/run-1/000001.tar").content + assert [entry["path"] for entry in listed] == ["changelog/run-1/000001.tar"] + assert increment == b"increment" diff --git a/uv.lock b/uv.lock index 79db1121..ee5bbe89 100644 --- a/uv.lock +++ b/uv.lock @@ -140,6 +140,7 @@ dependencies = [ [package.optional-dependencies] agent = [ { name = "a2a-sdk", extra = ["http-server"] }, + { name = "python-multipart" }, { name = "regex" }, { name = "uvicorn" }, ] @@ -148,6 +149,7 @@ dev = [ { name = "mcp" }, { name = "pytest" }, { name = "pytest-asyncio" }, + { name = "python-multipart" }, { name = "regex" }, ] @@ -160,6 +162,8 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.11" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=9.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0" }, + { name = "python-multipart", marker = "extra == 'agent'", specifier = ">=0.0.13" }, + { name = "python-multipart", marker = "extra == 'dev'", specifier = ">=0.0.13" }, { name = "regex", marker = "extra == 'agent'", specifier = ">=2023.12.25" }, { name = "regex", marker = "extra == 'dev'", specifier = ">=2023.12.25" }, { name = "starlette", specifier = ">=0.37" },