From 97a77700d758a0c757e5a0ee25d8d0e57af76ed2 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 06:50:51 -1000 Subject: [PATCH 01/14] fix(store): the local object store writes atomically and keeps content types A write now lands whole or not at all. Bytes are staged in a reserved .agentenv-tmp directory inside the root and put in place with a rename (an overwrite) or a hard link (a write-once put), so a crash or a concurrent reader never sees a partial object, and a write-once put that loses a race raises ObjectAlreadyExistsError instead of replacing the winner. Filesystems without hard links fall back to an exclusive create. A write's content type is recorded in a reserved .agentenv-meta tree and returned by get_object_metadata. The default type, which a writer that names none gets, reads back as unknown, so readers keep guessing from the name as before. Both directories are refused as keys and left out of listings. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 132 ++++++++++++++---- tst/unit/store/local_object_store_test.py | 90 +++++++++++- 2 files changed, 196 insertions(+), 26 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index 696726d6..5db87829 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -3,7 +3,12 @@ from __future__ import annotations import contextlib +import errno +import json +import logging +import os import shutil +import tempfile from collections.abc import Iterator from datetime import datetime, timezone from pathlib import Path @@ -13,11 +18,20 @@ from agent_env.store.local_state import ensure_state_dir from agent_env.store.object_store.object_store import DEFAULT_CONTENT_TYPE, ObjectMetadata, ObjectStore +logger = logging.getLogger(__name__) + +_META_DIR = ".agentenv-meta" # each object's content type, at the object's own key +_STAGING_DIR = ".agentenv-tmp" # files being written, inside the root so a rename into place is atomic +_RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) +# A filesystem without hard links: a no-overwrite write reserves its key by exclusive create instead. +_NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) + class LocalFilesystemObjectStore(ObjectStore): - """ObjectStore backed by a local directory tree; content types are not persisted. + """ObjectStore backed by a local directory tree. - The root-level ``.gitignore`` key is reserved for the ignore rule the store writes.""" + A write lands whole or not at all, and keeps any content type but the default. The root-level ``.gitignore``, + ``.agentenv-meta`` and ``.agentenv-tmp`` keys are reserved for the store's own files.""" def __init__(self, root: str) -> None: self._root = Path(root) @@ -36,18 +50,18 @@ def put( path = self._resolve(key) if path.exists() and not allow_overwrite: raise ObjectAlreadyExistsError(f"Object already exists at {path}.") - ensure_state_dir(self._root) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(data) + with self._staged() as staged: + staged.write_bytes(data) + self._commit(staged, path, content_type, allow_overwrite=allow_overwrite) return self._to_url(path) def put_file(self, key: str, file_path: str, content_type: str = DEFAULT_CONTENT_TYPE) -> str: path = self._resolve(key) if path.exists(): raise ObjectAlreadyExistsError(f"Object already exists at {path}.") - ensure_state_dir(self._root) - path.parent.mkdir(parents=True, exist_ok=True) - shutil.copyfile(file_path, path) + with self._staged() as staged: + shutil.copyfile(file_path, staged) + self._commit(staged, path, content_type, allow_overwrite=False) return self._to_url(path) def put_file_at(self, object_url: str, file_path: str, content_type: str = DEFAULT_CONTENT_TYPE) -> str: @@ -77,31 +91,25 @@ def get_object_metadata(self, key: str) -> ObjectMetadata | None: if not path.is_file(): return None st = path.stat() - return ObjectMetadata(content_type=None, size=st.st_size, last_modified=datetime.fromtimestamp(st.st_mtime, tz=timezone.utc)) + return ObjectMetadata( + content_type=self._read_content_type(path), + size=st.st_size, + last_modified=datetime.fromtimestamp(st.st_mtime, tz=timezone.utc), + ) def get_object_metadata_at(self, object_url: str) -> ObjectMetadata | None: return self.get_object_metadata(self.get_object_key(object_url)) def list(self, prefix: str) -> list[str]: - keys = [] - for path in self._root.rglob("*"): - if not path.is_file() or path == self._root / ".gitignore": - continue - key = path.relative_to(self._root).as_posix() - if key.startswith(prefix): - keys.append(key) - return keys + root = self._root.resolve() + return [key for key in (p.relative_to(root).as_posix() for p in self._objects_below(root)) if key.startswith(prefix)] def list_at(self, url_prefix: str) -> list[str]: base = url_prefix if url_prefix.endswith("/") else url_prefix + "/" dir_path = self._resolve(self.get_object_key(url_prefix)) if not dir_path.is_dir(): return [] - return [ - base + p.relative_to(dir_path).as_posix() - for p in sorted(dir_path.rglob("*")) - if p.is_file() and p != self._root.resolve() / ".gitignore" - ] + return [base + p.relative_to(dir_path).as_posix() for p in sorted(self._objects_below(dir_path))] def object_url(self, key: str) -> str: return self._to_url(self._resolve(key)) @@ -119,13 +127,87 @@ def get_object_key(self, object_url: str) -> str: return path.relative_to(root).as_posix() def _resolve(self, key: str) -> Path: + root = self._root.resolve() resolved = (self._root / key).resolve() - if not resolved.is_relative_to(self._root.resolve()): + if not resolved.is_relative_to(root): raise ValueError(f"Key {key!r} escapes the store root.") - if resolved == self._root.resolve() / ".gitignore": - raise ValueError(f"Key {key!r} is reserved: the store keeps its .gitignore there.") + parts = resolved.relative_to(root).parts + if parts and parts[0] in _RESERVED: + raise ValueError(f"Key {key!r} is reserved: the store keeps its own {parts[0]} there.") return resolved + def _objects_below(self, dir_path: Path) -> Iterator[Path]: + """Every object file below ``dir_path``, leaving out the store's own files at the root.""" + root = self._root.resolve() + if not dir_path.is_dir(): + return + for entry in dir_path.iterdir(): + if dir_path == root and entry.name in _RESERVED: + continue + if entry.is_file(): + yield entry + elif entry.is_dir(): + yield from (p for p in entry.rglob("*") if p.is_file()) + + @contextlib.contextmanager + def _staged(self) -> Iterator[Path]: + """An empty file to write an object into before it is committed; removed if it never is.""" + ensure_state_dir(self._root) + staging = self._root / _STAGING_DIR + staging.mkdir(exist_ok=True) + fd, name = tempfile.mkstemp(dir=staging) + os.close(fd) + staged = Path(name) + staged.chmod(0o644) + try: + yield staged + finally: + staged.unlink(missing_ok=True) + + def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrite: bool) -> None: + """Put a staged file in place at ``path`` and record its content type. An overwrite renames over + whatever is there; otherwise a hard link claims ``path`` only if no other writer got there first.""" + path.parent.mkdir(parents=True, exist_ok=True) + if allow_overwrite: + os.replace(staged, path) + else: + try: + os.link(staged, path) + except FileExistsError: + raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None + except OSError as e: + if e.errno not in _NO_HARD_LINKS: + raise + try: + os.close(os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)) + except FileExistsError: + raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None + os.replace(staged, path) + # The default type is what a writer that named none gets; it reads back as unknown, so readers + # still guess from the name as they did before types were kept. + try: + meta = self._meta_path(path) + if content_type == DEFAULT_CONTENT_TYPE: + meta.unlink(missing_ok=True) + return + meta.parent.mkdir(parents=True, exist_ok=True) + with self._staged() as staged_meta: + staged_meta.write_text(json.dumps({"content_type": content_type})) + os.replace(staged_meta, meta) + except OSError as e: # the object is in place; only its type goes unrecorded + logger.warning("Could not record the content type of %s: %s", path, e) + + def _read_content_type(self, path: Path) -> str | None: + """The content type recorded for the object at ``path``; None for one written before types were kept.""" + try: + return json.loads(self._meta_path(path).read_text()).get("content_type") + except (OSError, ValueError, AttributeError): + return None + + def _meta_path(self, path: Path) -> Path: + root = self._root.resolve() + return root / _META_DIR / path.relative_to(root) + @staticmethod def _to_url(path: Path) -> str: return f"file://{path.resolve()}" diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index aec1fe12..3ac8e048 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -4,9 +4,13 @@ assertions that prove S3 parity also give quick backend-neutral coverage. """ +import errno +import os +from pathlib import Path + import pytest -from agent_env.store import LocalFilesystemObjectStore +from agent_env.store import LocalFilesystemObjectStore, ObjectAlreadyExistsError from tst.store import object_conformance @@ -105,3 +109,87 @@ def test_signed_get_url_is_none(store): def test_signed_put_url_is_none(store): """Local FS can't presign uploads either — VM-upload flows require a signable backend.""" assert store.signed_put_url("file:///whatever") is None + + +def test_a_write_keeps_its_content_type(store): + store.put("typed/a.json", b"{}", content_type="application/json") + store.put_file("typed/b.csv", __file__, content_type="text/csv") + assert store.get_object_metadata("typed/a.json").content_type == "application/json" + assert store.get_object_metadata("typed/b.csv").content_type == "text/csv" + + +def test_a_write_that_names_no_type_reads_back_as_unknown(store): + """Readers then guess from the name, as they did before types were kept.""" + store.put("untyped/a.zip", b"PK") + assert store.get_object_metadata("untyped/a.zip").content_type is None + + +def test_an_overwrite_replaces_the_recorded_type(store): + store.put("over/x", b"1", content_type="application/json") + store.put("over/x", b"2", content_type="text/plain", allow_overwrite=True) + assert store.get_object_metadata("over/x").content_type == "text/plain" + store.put("over/x", b"3", allow_overwrite=True) + assert store.get_object_metadata("over/x").content_type is None + + +def test_an_object_written_before_types_were_kept_has_none(store, tmp_path): + (tmp_path / "old").mkdir() + (tmp_path / "old" / "x.bin").write_bytes(b"legacy") + md = store.get_object_metadata("old/x.bin") + assert (md.content_type, md.size) == (None, 6) + + +def test_a_failed_write_leaves_no_object(store, monkeypatch): + def fail(self, data): + raise OSError("disk full") + + monkeypatch.setattr(Path, "write_bytes", fail) + with pytest.raises(OSError, match="disk full"): + store.put("partial/x.bin", b"payload") + monkeypatch.undo() + assert not store.exists("partial/x.bin") and store.list("") == [] + + +def test_a_no_overwrite_write_that_loses_a_race_keeps_the_winner(store, monkeypatch): + """Another writer lands the key between the existence check and the commit.""" + real_link = os.link + + def racing_link(src, dst): + Path(dst).write_bytes(b"winner") + return real_link(src, dst) + + monkeypatch.setattr(os, "link", racing_link) + with pytest.raises(ObjectAlreadyExistsError): + store.put("race/x", b"loser") + assert store.read("race/x") == b"winner" + + +def test_a_filesystem_without_hard_links_still_refuses_an_existing_key(store, monkeypatch): + def no_links(src, dst): + raise OSError(errno.EPERM, "hard links not supported") + + monkeypatch.setattr(os, "link", no_links) + url = store.put("nolink/x", b"first", content_type="text/plain") + assert store.get(url) == b"first" + with pytest.raises(ObjectAlreadyExistsError): + store.put("nolink/x", b"second") + assert store.read("nolink/x") == b"first" + + +@pytest.mark.parametrize("key", [".agentenv-meta/x", ".agentenv-tmp/x", "a/../.agentenv-meta"]) +def test_the_stores_own_directories_are_reserved(store, key): + with pytest.raises(ValueError, match="reserved"): + store.put(key, b"mine") + + +def test_listings_leave_out_the_stores_own_files(store, tmp_path): + store.put("a/x.json", b"{}", content_type="application/json") + (tmp_path / ".agentenv-tmp").mkdir(exist_ok=True) + (tmp_path / ".agentenv-tmp" / "crash-residue").write_bytes(b"half") + assert store.list("") == ["a/x.json"] + assert store.list_at(store.object_url("")) == [store.object_url("a/x.json")] + assert os.listdir(tmp_path / ".agentenv-tmp") == ["crash-residue"] # committed writes leave nothing staged + + +def test_list_of_an_unwritten_store_is_empty(tmp_path): + assert LocalFilesystemObjectStore(str(tmp_path / "never-written")).list("") == [] From 8ef1ddf0a7015d679e0c18617390352f34365bd5 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 07:15:37 -1000 Subject: [PATCH 02/14] feat(store): the local object store issues HTTPS transfer grants LocalFilesystemObjectStore's issue_read_grant, issue_write_grant and issue_upload_policy now return real http-get, http-put and http-post-policy grants, served by an HTTPS server inside the process that issued them. It is started by the first grant and runs in a daemon thread: starlette on uvicorn, one route, /v1/grants/{token}, with no access log. - Tokens are HMAC-SHA256-signed claims (operation, key or key prefix, expiry, size bound, content type) under a key held only in this process's memory; the server trusts nothing else. - GET returns the stored bytes and type. PUT requires the signed Content-Type, bounds the body to the grant while it streams, and replaces the object atomically. A multipart POST names its key before its file, must stay under the grant's prefix, and is bounded the same way. - TLS comes from a per-user local CA under /tls, name-constrained to localhost, host.docker.internal and loopback or private addresses, with ca.pem and ca-bundle.pem (the public roots plus the CA) for containers to trust. It is never installed on the host. - The server listens on 127.0.0.1, or the Docker bridge gateway on Linux, and grant URLs name host.docker.internal; grant_bind_host and grant_advertise_host override both. root is now optional and defaults to /object_store. supports_transfer_grants stays false, so no agent is handed these grants yet: containers must first trust the CA and reach the server. starlette, uvicorn (0.47 for ssl_context_factory), python-multipart, cryptography and certifi are declared as core dependencies; all were already installed. Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 5 + .../object_store/local_grants/__init__.py | 2 + .../store/object_store/local_grants/server.py | 386 ++++++++++++++++++ .../store/object_store/local_grants/tls.py | 230 +++++++++++ .../store/object_store/local_grants/tokens.py | 119 ++++++ .../store/object_store/local_object_store.py | 74 +++- tst/unit/store/local_grants_server_test.py | 270 ++++++++++++ tst/unit/store/local_grants_tls_test.py | 107 +++++ tst/unit/store/local_grants_tokens_test.py | 65 +++ tst/unit/store/test_config_object_store.py | 14 + tst/unit/store/test_object_transfer_grants.py | 11 +- uv.lock | 10 + 12 files changed, 1279 insertions(+), 14 deletions(-) create mode 100644 src/agent_env/store/object_store/local_grants/__init__.py create mode 100644 src/agent_env/store/object_store/local_grants/server.py create mode 100644 src/agent_env/store/object_store/local_grants/tls.py create mode 100644 src/agent_env/store/object_store/local_grants/tokens.py create mode 100644 tst/unit/store/local_grants_server_test.py create mode 100644 tst/unit/store/local_grants_tls_test.py create mode 100644 tst/unit/store/local_grants_tokens_test.py diff --git a/pyproject.toml b/pyproject.toml index ea3c9f21..b180b724 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,9 @@ requires-python = ">=3.11" dependencies = [ "agentenv-protocol>=0.1.0", "boto3>=1.35.0", + "certifi>=2024.2.2", # the public roots the local transfer CA's trust bundle carries "click>=8.0.0", + "cryptography>=42", # the local transfer CA "httpx>=0.27.0", "litellm>=1.96.0,<1.97", "mcp>=1.0.0,<2", # 2.0.0 removed mcp.server.fastmcp, imported at module scope by env/gateway @@ -22,7 +24,10 @@ dependencies = [ "a2a-sdk==0.3.26", "pydantic>=2.12", # model_validate(extra=...) reads agent responses tolerantly "pymongo>=4.7", + "python-multipart>=0.0.13", # the local grant server's upload forms "pyyaml>=6.0", + "starlette>=0.40", # the local grant server + "uvicorn>=0.47", # 0.47 adds ssl_context_factory, how the local grant server gets its TLS context "modal>=1.5.2.dev2", # dev2 adds VM-mode V2 (_experimental_create + vm_runtime) "e2b~=2.46.4", # AsyncSandbox/AsyncTemplate APIs validated against the 2.46 line ] diff --git a/src/agent_env/store/object_store/local_grants/__init__.py b/src/agent_env/store/object_store/local_grants/__init__.py new file mode 100644 index 00000000..4fdd242f --- /dev/null +++ b/src/agent_env/store/object_store/local_grants/__init__.py @@ -0,0 +1,2 @@ +"""HTTPS object-transfer grants for the local filesystem object store: signed tokens, a local +certificate authority, and the in-process server that honours them.""" diff --git a/src/agent_env/store/object_store/local_grants/server.py b/src/agent_env/store/object_store/local_grants/server.py new file mode 100644 index 00000000..cf2449ce --- /dev/null +++ b/src/agent_env/store/object_store/local_grants/server.py @@ -0,0 +1,386 @@ +"""The in-process HTTPS server that honours local object-transfer grants. + +One route, ``/v1/grants/{token}``, with the semantics of the S3 requests the grants stand in for: +GET returns an object's exact bytes; PUT writes one object of the claimed type and size; a +multipart POST writes one object anywhere under the claimed prefix. A token that does not verify, +has expired or names another method is a 403; a request outside the grant's bounds is a 4xx. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import functools +import ipaddress +import logging +import platform +import socket +import subprocess +import threading +import time +from pathlib import Path +from typing import TYPE_CHECKING, BinaryIO + +import uvicorn +from python_multipart.exceptions import FormParserError +from python_multipart.multipart import MultipartParser, parse_options_header +from starlette.applications import Starlette +from starlette.requests import ClientDisconnect, Request +from starlette.responses import FileResponse, PlainTextResponse, Response +from starlette.routing import Route +from uvicorn.protocols.http.h11_impl import H11Protocol + +from agent_env.store.base import GrantUnavailableError +from agent_env.store.object_store.local_grants.tls import check_local_host, local_ca, server_context +from agent_env.store.object_store.local_grants.tokens import ( + GrantClaims, + GrantSigner, + InvalidGrantError, + Op, + store_id, +) +from agent_env.store.object_store.object_store import DEFAULT_CONTENT_TYPE + +if TYPE_CHECKING: + from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore + +logger = logging.getLogger(__name__) + +DEFAULT_ADVERTISE_HOST = "host.docker.internal" +_START_TIMEOUT_SECONDS = 10 +_OPS_BY_METHOD = {"GET": "get", "HEAD": "get", "PUT": "put", "POST": "post"} +_UPLOAD_KEY_FIELD = "key" # the form fields a policy upload names its key and file by, as S3's do +_UPLOAD_FILE_FIELD = "file" +_MAX_FORM_FIELDS = 32 +_MAX_FORM_FIELD_BYTES = 64 * 1024 + +_servers: dict[tuple[str, str], GrantServer] = {} +_servers_lock = threading.Lock() + + +class _UnloggedH11Protocol(H11Protocol): + """HTTP/1.1 without access logs, which would record grant tokens. uvicorn's own switch clears + the process-wide access logger, silencing any other server in this process.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.access_log = False + + +class _Rejected(Exception): + def __init__(self, status: int, reason: str) -> None: + super().__init__(reason) + self.status = status + self.reason = reason + + +def grant_server(bind_host: str | None, advertise_host: str | None) -> GrantServer: + """This process's grant server for a bind and advertised host, created on first use; neither + given means this platform's default bind address and ``host.docker.internal``.""" + key = (bind_host or default_bind_host(), advertise_host or DEFAULT_ADVERTISE_HOST) + with _servers_lock: + server = _servers.get(key) + if server is None: + server = _servers[key] = GrantServer(*key) + return server + + +@functools.cache +def default_bind_host() -> str: + """Where the grant server listens unless told otherwise. On macOS and Windows, Docker reaches + the host's loopback as ``host.docker.internal``; native Linux maps that name to the default + bridge's gateway (``--add-host host.docker.internal:host-gateway``), so the server listens + there, reachable from containers and not beyond the host.""" + if platform.system() != "Linux": + return "127.0.0.1" + try: + gateway = subprocess.run( + ["docker", "network", "inspect", "bridge", "--format", "{{range .IPAM.Config}}{{.Gateway}} {{end}}"], + capture_output=True, text=True, timeout=10, check=True, + ).stdout.split() + return str(ipaddress.ip_address(gateway[0])) + except (OSError, subprocess.SubprocessError, IndexError, ValueError) as e: + logger.warning( + "Could not find the Docker bridge gateway (%s); the local grant server listens on 127.0.0.1, " + "which containers may not reach. Set grant_bind_host under [stores.object] to override.", e, + ) + return "127.0.0.1" + + +class GrantServer: + """Serves the grants of the local stores registered with it, over HTTPS on one port. Started by + the first grant it issues; it runs in a daemon thread for the rest of the process.""" + + def __init__(self, bind_host: str, advertise_host: str) -> None: + self.bind_host = bind_host + self.advertise_host = advertise_host + self._signer = GrantSigner() + self._stores: dict[str, LocalFilesystemObjectStore] = {} + self._lock = threading.Lock() + self._port: int | None = None + self._uvicorn: uvicorn.Server | None = None + self._thread: threading.Thread | None = None + self._app = Starlette(routes=[Route("/v1/grants/{token}", self._handle, methods=["GET", "PUT", "POST"])]) + + def issue( + self, + store: LocalFilesystemObjectStore, + op: Op, + expires: int, + *, + key: str | None = None, + prefix: str | None = None, + max_bytes: int | None = None, + content_type: str | None = None, + ) -> str: + """The URL of a new grant for ``op`` on ``store``, valid until ``expires`` (epoch seconds).""" + claims = GrantClaims( + op=op, store=self._register(store), expires=expires, key=key, prefix=prefix, + max_bytes=max_bytes, content_type=content_type, + ) + port = self._start() + host = f"[{self.advertise_host}]" if ":" in self.advertise_host else self.advertise_host + return f"https://{host}:{port}/v1/grants/{self._signer.sign(claims)}" + + def close(self) -> None: + """Stop serving; grants already issued stop working.""" + with self._lock: + if self._uvicorn is not None: + self._uvicorn.should_exit = True + if self._thread is not None: + self._thread.join(timeout=_START_TIMEOUT_SECONDS) + self._port = self._uvicorn = self._thread = None + + def _register(self, store: LocalFilesystemObjectStore) -> str: + name = store_id(str(store.root.resolve())) + with self._lock: + self._stores[name] = store + return name + + def _start(self) -> int: + with self._lock: + if self._port is not None: + return self._port + address = ipaddress.ip_address(self.bind_host) + sock = socket.socket(socket.AF_INET6 if address.version == 6 else socket.AF_INET, socket.SOCK_STREAM) + try: + sock.bind((self.bind_host, 0)) + hosts = ["localhost", "host.docker.internal", "127.0.0.1", self.advertise_host] + with contextlib.suppress(ValueError): # a name the CA cannot vouch for would void the certificate + check_local_host(self.bind_host) + hosts.append(self.bind_host) + context = server_context(local_ca(), hosts) + except (OSError, ValueError) as e: + sock.close() + raise GrantUnavailableError(f"the local grant server cannot listen on {self.bind_host}: {e}") from e + config = uvicorn.Config( + self._app, lifespan="off", log_config=None, http=_UnloggedH11Protocol, + ssl_context_factory=lambda _config, _default: context, + ) + server = uvicorn.Server(config) + failure: list[BaseException] = [] + + def serve() -> None: + try: + asyncio.run(server.serve(sockets=[sock])) + except BaseException as e: # uvicorn ends a failed startup with SystemExit + failure.append(e) + + thread = threading.Thread(target=serve, name="agent-env-grant-server", daemon=True) + thread.start() + deadline = time.monotonic() + _START_TIMEOUT_SECONDS + while not server.started: + if failure or not thread.is_alive() or time.monotonic() > deadline: + server.should_exit = True + sock.close() + raise GrantUnavailableError( + f"the local grant server did not start on {self.bind_host}" + + (f": {failure[0]!r}" if failure else "") + ) + time.sleep(0.01) + self._uvicorn, self._thread = server, thread + self._port = sock.getsockname()[1] + logger.info("Local grant server listening on %s:%d", self.bind_host, self._port) + return self._port + + async def _handle(self, request: Request) -> Response: + try: + claims = self._signer.verify(request.path_params["token"], now=time.time()) + except (InvalidGrantError, ValueError, KeyError, TypeError): + return PlainTextResponse("The grant is invalid or has expired.", status_code=403) + store = self._stores.get(claims.store) + if store is None or claims.op != _OPS_BY_METHOD.get(request.method): + return PlainTextResponse("The grant does not allow this request.", status_code=403) + try: + if claims.op == "get": + return _get(store, claims) + if claims.op == "put": + await _put(store, claims, request) + return Response(status_code=200) + await _post(store, claims, request) + return Response(status_code=204) + except _Rejected as r: + return PlainTextResponse(r.reason, status_code=r.status) + except ClientDisconnect: + return PlainTextResponse("The upload ended early.", status_code=400) + except Exception: + logger.exception("The local grant server failed a %s request", request.method) + return PlainTextResponse("The object store failed the transfer.", status_code=500) + + +def _path(store: LocalFilesystemObjectStore, key: str) -> Path: + try: + return store._resolve(key) + except ValueError: + raise _Rejected(403, "The grant's key is not an object key of this store.") from None + + +def _get(store: LocalFilesystemObjectStore, claims: GrantClaims) -> Response: + path = _path(store, claims.key) + if not path.is_file(): + raise _Rejected(404, "No object exists at this grant's key.") + return FileResponse(path, media_type=store._read_content_type(path) or DEFAULT_CONTENT_TYPE) + + +async def _put(store: LocalFilesystemObjectStore, claims: GrantClaims, request: Request) -> None: + if request.headers.get("content-type") != claims.content_type: + raise _Rejected(403, f"The upload must declare Content-Type {claims.content_type}.") + declared = request.headers.get("content-length") + try: + if declared is not None and int(declared) > claims.max_bytes: + raise _Rejected(400, "The upload is larger than the grant allows.") + except ValueError: + raise _Rejected(400, "The upload's Content-Length is not a number.") from None + path = _path(store, claims.key) + with store._staged() as staged: + received = 0 + with staged.open("wb") as out: + async for chunk in request.stream(): + received += len(chunk) + if received > claims.max_bytes: + raise _Rejected(400, "The upload is larger than the grant allows.") + out.write(chunk) + store._commit(staged, path, claims.content_type, allow_overwrite=True) + + +async def _post(store: LocalFilesystemObjectStore, claims: GrantClaims, request: Request) -> None: + 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 _Rejected(400, "The upload must be multipart/form-data.") + with contextlib.ExitStack() as stack: + upload = _PolicyUpload(store, claims, stack) + parser = MultipartParser(params[b"boundary"], upload.callbacks()) + try: + async for chunk in request.stream(): + parser.write(chunk) + parser.finalize() + except FormParserError: + raise _Rejected(400, "The upload is not a valid multipart form.") from None + upload.commit() + + +class _PolicyUpload: + """One multipart POST under a policy grant, read as S3 reads one: form fields, then the file, + which is written to a staged file as it arrives. Fields after the file are ignored.""" + + def __init__(self, store: LocalFilesystemObjectStore, claims: GrantClaims, stack: contextlib.ExitStack) -> None: + self._store = store + self._claims = claims + self._stack = stack + self._fields: dict[str, bytes] = {} + self._headers: dict[bytes, bytes] = {} + self._header_name = b"" + self._header_value = b"" + self._part_name: str | None = None + self._part_data = bytearray() + self._path: Path | None = None + self._staged: Path | None = None + self._out: BinaryIO | None = None + self._received = 0 + 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 commit(self) -> None: + if self._staged is None or not self._file_done: + raise _Rejected(400, f"The upload has no {_UPLOAD_FILE_FIELD!r} field.") + self._out.close() + content_type = self._field("Content-Type") or DEFAULT_CONTENT_TYPE + self._store._commit(self._staged, self._path, content_type, allow_overwrite=True) + + def _on_part_begin(self) -> None: + self._headers = {} + self._part_name = None + self._part_data = bytearray() + + def _on_header_field(self, data: bytes, start: int, end: int) -> None: + self._header_name += data[start:end] + + def _on_header_value(self, data: bytes, start: int, end: int) -> None: + self._header_value += data[start:end] + + 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 _Rejected(400, "A form field has no name.") + self._part_name = name.decode("utf-8", "replace") + if self._part_name == _UPLOAD_FILE_FIELD: + self._path = _path(self._store, self._upload_key()) + self._staged = self._stack.enter_context(self._store._staged()) + self._out = self._stack.enter_context(self._staged.open("wb")) + elif len(self._fields) >= _MAX_FORM_FIELDS: + raise _Rejected(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 == _UPLOAD_FILE_FIELD: + self._received += end - start + if self._received > self._claims.max_bytes: + raise _Rejected(400, "The upload is larger than the grant allows.") + self._out.write(data[start:end]) + return + if len(self._part_data) + end - start > _MAX_FORM_FIELD_BYTES: + raise _Rejected(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 == _UPLOAD_FILE_FIELD: + self._file_done = True + elif self._part_name is not None: + self._fields[self._part_name] = bytes(self._part_data) + + def _field(self, name: str) -> str | None: + value = self._fields.get(name) + return None if value is None else value.decode("utf-8", "replace") + + def _upload_key(self) -> str: + """The key the upload names, which must come before the file, be a plain relative path, + and lie under the grant's prefix.""" + key = self._field(_UPLOAD_KEY_FIELD) + if key is None: + raise _Rejected(400, f"The upload must name its {_UPLOAD_KEY_FIELD!r} before its file.") + parts = key.split("/") + if key.startswith("/") or "\\" in key or "\x00" in key or any(p in ("", ".", "..") for p in parts): + raise _Rejected(403, "The upload key is not a normalized relative path.") + if not key.startswith(self._claims.prefix): + raise _Rejected(403, "The upload key is outside the grant's prefix.") + return key diff --git a/src/agent_env/store/object_store/local_grants/tls.py b/src/agent_env/store/object_store/local_grants/tls.py new file mode 100644 index 00000000..d595db5d --- /dev/null +++ b/src/agent_env/store/object_store/local_grants/tls.py @@ -0,0 +1,230 @@ +"""The local certificate authority behind the grant server's HTTPS. + +One CA per user, created on first use under ``/tls/``. It is name-constrained to +local names and private addresses, so even a leaked key can only vouch for those, and it is +never installed in the host's trust stores: only the containers agent-env starts trust it, +through ``ca-bundle.pem`` (the public roots plus this CA, for ``SSL_CERT_FILE``, which replaces +the defaults) or ``ca.pem`` (the CA alone, for trust stores that add to the defaults). +""" + +from __future__ import annotations + +import ipaddress +import os +import secrets +import ssl +import tempfile +from collections.abc import Iterable +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from pathlib import Path + +import certifi +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID + +from agent_env.config.paths import state_root +from agent_env.store.local_state import ensure_state_dir + +# What the CA may vouch for: the names containers reach the host by, and loopback and private addresses. +LOCAL_NAMES = ("localhost", "host.docker.internal") +LOCAL_NETWORKS = tuple( + ipaddress.ip_network(n) + for n in ("127.0.0.0/8", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "::1/128", "fc00::/7") +) + +_CA_LIFETIME = timedelta(days=5 * 365) +_CA_RENEW_BEFORE = timedelta(days=30) # replaced this long before it expires +_SERVER_LIFETIME = timedelta(days=90) +_CLOCK_SKEW = timedelta(minutes=5) + +_CA_FILE = "ca.key.pem" # the CA's key and certificate, private to the user +_CERT_FILE = "ca.pem" +_BUNDLE_FILE = "ca-bundle.pem" +_PEM_CERT = b"-----BEGIN CERTIFICATE-----" + + +@dataclass(frozen=True) +class LocalCA: + """The CA's key and certificate, and the trust files derived from it.""" + + key: ec.EllipticCurvePrivateKey + cert: x509.Certificate + cert_path: Path + bundle_path: Path + + +def check_local_host(host: str) -> None: + """Refuse a host the CA cannot vouch for: a name other than the local ones, or a public address.""" + try: + address = ipaddress.ip_address(host) + except ValueError: + name = host.lower().rstrip(".") + if any(name == n or name.endswith("." + n) for n in LOCAL_NAMES): + return + raise ValueError( + f"{host!r} is not a local name; the local transfer CA vouches only for {', '.join(LOCAL_NAMES)} " + "and loopback or private addresses" + ) from None + if not any(address in network for network in LOCAL_NETWORKS): + raise ValueError(f"{host!r} is not a loopback or private address, which the local transfer CA vouches for") + + +def local_ca() -> LocalCA: + """This user's local CA, created on first use and replaced when it nears expiry.""" + tls_dir = state_root() / "tls" + ensure_state_dir(tls_dir) + ca_file = tls_dir / _CA_FILE + loaded = _load(ca_file) + if loaded is None: + key, cert = _new_ca() + pem = _key_pem(key) + cert.public_bytes(serialization.Encoding.PEM) + # A first CA is claimed with a hard link, so two processes starting at once agree on one. + replacing = ca_file.exists() + staged = _write_private(tls_dir, pem) + try: + if replacing: + os.replace(staged, ca_file) + else: + try: + os.link(staged, ca_file) + except FileExistsError: + pass + finally: + staged.unlink(missing_ok=True) + loaded = _load(ca_file) or (key, cert) + key, cert = loaded + cert_pem = cert.public_bytes(serialization.Encoding.PEM) + cert_path = _write_if_changed(tls_dir / _CERT_FILE, cert_pem) + bundle_path = _write_if_changed(tls_dir / _BUNDLE_FILE, Path(certifi.where()).read_bytes().rstrip() + b"\n" + cert_pem) + return LocalCA(key=key, cert=cert, cert_path=cert_path, bundle_path=bundle_path) + + +def server_context(ca: LocalCA, hosts: Iterable[str]) -> ssl.SSLContext: + """A TLS server context presenting a new certificate, signed by ``ca``, for each of ``hosts``.""" + key = ec.generate_private_key(ec.SECP256R1()) + names = list(dict.fromkeys(hosts)) + now = datetime.now(UTC) + cert = ( + x509.CertificateBuilder() + .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")])) + .issuer_name(ca.cert.subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - _CLOCK_SKEW) + .not_valid_after(min(now + _SERVER_LIFETIME, ca.cert.not_valid_after_utc)) + .add_extension(x509.SubjectAlternativeName([_general_name(n) for n in names]), critical=False) + .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True) + .add_extension(_key_usage(digital_signature=True), critical=True) + .add_extension(x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]), critical=False) + .add_extension(x509.SubjectKeyIdentifier.from_public_key(key.public_key()), critical=False) + .add_extension( + x509.AuthorityKeyIdentifier.from_issuer_public_key(ca.key.public_key()), critical=False + ) + .sign(ca.key, hashes.SHA256()) + ) + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.minimum_version = ssl.TLSVersion.TLSv1_2 + # ssl loads a certificate chain only from a file; this one lives just long enough to be read. + with tempfile.TemporaryDirectory() as tmp: + chain = Path(tmp) / "server.pem" + chain.write_bytes(_key_pem(key) + cert.public_bytes(serialization.Encoding.PEM)) + context.load_cert_chain(chain) + return context + + +def _new_ca() -> tuple[ec.EllipticCurvePrivateKey, x509.Certificate]: + key = ec.generate_private_key(ec.SECP256R1()) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, f"agent-env local transfer CA {secrets.token_hex(4)}")]) + now = datetime.now(UTC) + cert = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - _CLOCK_SKEW) + .not_valid_after(now + _CA_LIFETIME) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .add_extension(_key_usage(key_cert_sign=True, crl_sign=True), critical=True) + .add_extension(_name_constraints(), critical=True) + .add_extension(x509.SubjectKeyIdentifier.from_public_key(key.public_key()), critical=False) + .sign(key, hashes.SHA256()) + ) + return key, cert + + +def _load(ca_file: Path) -> tuple[ec.EllipticCurvePrivateKey, x509.Certificate] | None: + """The CA in ``ca_file``, or None when it is missing, unreadable, nearly expired, or not one this + release would create.""" + try: + data = ca_file.read_bytes() + key_pem, marker, cert_pem = data.partition(_PEM_CERT) + key = serialization.load_pem_private_key(key_pem, password=None) + cert = x509.load_pem_x509_certificate(marker + cert_pem) + constraints = cert.extensions.get_extension_for_class(x509.NameConstraints).value + except (OSError, ValueError, TypeError, x509.ExtensionNotFound): + return None + if not isinstance(key, ec.EllipticCurvePrivateKey) or key.public_key() != cert.public_key(): + return None + if constraints != _name_constraints() or cert.not_valid_after_utc - datetime.now(UTC) < _CA_RENEW_BEFORE: + return None + return key, cert + + +def _name_constraints() -> x509.NameConstraints: + return x509.NameConstraints( + permitted_subtrees=[x509.DNSName(n) for n in LOCAL_NAMES] + [x509.IPAddress(n) for n in LOCAL_NETWORKS], + excluded_subtrees=None, + ) + + +def _key_usage(*, digital_signature: bool = False, key_cert_sign: bool = False, crl_sign: bool = False) -> x509.KeyUsage: + return x509.KeyUsage( + digital_signature=digital_signature, + content_commitment=False, + key_encipherment=False, + data_encipherment=False, + key_agreement=False, + key_cert_sign=key_cert_sign, + crl_sign=crl_sign, + encipher_only=False, + decipher_only=False, + ) + + +def _general_name(host: str) -> x509.GeneralName: + try: + return x509.IPAddress(ipaddress.ip_address(host)) + except ValueError: + return x509.DNSName(host) + + +def _key_pem(key: ec.EllipticCurvePrivateKey) -> bytes: + return key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ) + + +def _write_private(directory: Path, data: bytes) -> Path: + fd, name = tempfile.mkstemp(dir=directory, prefix=".ca.") # created 0600 + with os.fdopen(fd, "wb") as f: + f.write(data) + return Path(name) + + +def _write_if_changed(path: Path, data: bytes) -> Path: + try: + if path.read_bytes() == data: + return path + except OSError: + pass + staged = _write_private(path.parent, data) + try: + staged.chmod(0o644) + os.replace(staged, path) + finally: + staged.unlink(missing_ok=True) + return path diff --git a/src/agent_env/store/object_store/local_grants/tokens.py b/src/agent_env/store/object_store/local_grants/tokens.py new file mode 100644 index 00000000..2af0c3cb --- /dev/null +++ b/src/agent_env/store/object_store/local_grants/tokens.py @@ -0,0 +1,119 @@ +"""Grant tokens: signed claims naming one operation on one key or key prefix, until an expiry. + +A token is ``v1..``: base64url JSON claims and their HMAC-SHA256 under a key +that lives only in this process's memory, so a token stops working when the process exits and +the server trusts nothing but the claims it signed. +""" + +from __future__ import annotations + +import base64 +import binascii +import hashlib +import hmac +import json +import secrets +from dataclasses import dataclass, field +from typing import Literal + +_VERSION = "v1" +_OPS = ("get", "put", "post") + +Op = Literal["get", "put", "post"] + + +class InvalidGrantError(Exception): + """A token this process did not sign, or one that has expired.""" + + +@dataclass(frozen=True) +class GrantClaims: + """What a grant allows. ``key`` names the object of a get or put; ``prefix``, ending in ``/``, + the keys a post may write. ``max_bytes`` bounds an upload and ``content_type`` is the type a + put must declare.""" + + op: Op + store: str + expires: int + key: str | None = None + prefix: str | None = None + max_bytes: int | None = None + content_type: str | None = None + grant_id: str = field(default_factory=lambda: secrets.token_hex(8)) + + def __post_init__(self) -> None: + if self.op not in _OPS: + raise ValueError(f"op must be one of {_OPS}, got {self.op!r}") + if self.op == "post": + if self.key is not None or self.prefix is None or not (self.prefix == "" or self.prefix.endswith("/")): + raise ValueError("a post grant names a prefix ending in '/' (or the empty root prefix), not a key") + elif self.key is None or self.prefix is not None: + raise ValueError(f"a {self.op} grant names one key, not a prefix") + if self.op != "get" and (not isinstance(self.max_bytes, int) or self.max_bytes < 0): + raise ValueError(f"a {self.op} grant needs a non-negative max_bytes") + if (self.op == "put") != (self.content_type is not None): + raise ValueError("a put grant, and only a put grant, names a content type") + + +class GrantSigner: + """Signs and verifies grant tokens with a key generated for this process.""" + + def __init__(self) -> None: + self._key = secrets.token_bytes(32) + + def sign(self, claims: GrantClaims) -> str: + body = { + "op": claims.op, + "s": claims.store, + "exp": claims.expires, + "id": claims.grant_id, + "k": claims.key, + "p": claims.prefix, + "max": claims.max_bytes, + "ct": claims.content_type, + } + payload = _b64encode(json.dumps({k: v for k, v in body.items() if v is not None}, separators=(",", ":")).encode()) + return f"{_VERSION}.{payload}.{_b64encode(self._mac(payload))}" + + def verify(self, token: str, *, now: float) -> GrantClaims: + """The claims of ``token`` if this signer signed it and it has not expired at ``now``.""" + version, _, rest = token.partition(".") + payload, _, signature = rest.partition(".") + if version != _VERSION or not payload or not signature: + raise InvalidGrantError("not a grant token") + try: + valid = hmac.compare_digest(_b64decode(signature), self._mac(payload)) + except (binascii.Error, ValueError): + valid = False + if not valid: + raise InvalidGrantError("the grant's signature does not match") + body = json.loads(_b64decode(payload)) + claims = GrantClaims( + op=body["op"], + store=body["s"], + expires=body["exp"], + grant_id=body["id"], + key=body.get("k"), + prefix=body.get("p"), + max_bytes=body.get("max"), + content_type=body.get("ct"), + ) + if claims.expires <= now: + raise InvalidGrantError("the grant has expired") + return claims + + def _mac(self, payload: str) -> bytes: + return hmac.new(self._key, f"{_VERSION}.{payload}".encode(), hashlib.sha256).digest() + + +def store_id(root: str) -> str: + """A short, stable name for the store at ``root``, so a token names its store without its path.""" + return hashlib.sha256(root.encode()).hexdigest()[:16] + + +def _b64encode(data: bytes) -> str: + return base64.urlsafe_b64encode(data).rstrip(b"=").decode() + + +def _b64decode(text: str) -> bytes: + return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4)) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index 5db87829..de9bb521 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -4,19 +4,32 @@ import contextlib import errno +import ipaddress import json import logging import os import shutil import tempfile +import time from collections.abc import Iterator from datetime import datetime, timezone from pathlib import Path from typing import BinaryIO +from agentenv_protocol.transfers import HttpGetGrant, HttpPostPolicyGrant, HttpPutGrant + +from agent_env.config.paths import state_root from agent_env.store.base import ObjectAlreadyExistsError, ObjectNotFoundError from agent_env.store.local_state import ensure_state_dir -from agent_env.store.object_store.object_store import DEFAULT_CONTENT_TYPE, ObjectMetadata, ObjectStore +from agent_env.store.object_store.local_grants.server import grant_server +from agent_env.store.object_store.local_grants.tls import check_local_host +from agent_env.store.object_store.local_grants.tokens import Op +from agent_env.store.object_store.object_store import ( + DEFAULT_CONTENT_TYPE, + ObjectMetadata, + ObjectStore, + UploadPolicy, +) logger = logging.getLogger(__name__) @@ -28,13 +41,29 @@ class LocalFilesystemObjectStore(ObjectStore): - """ObjectStore backed by a local directory tree. + """ObjectStore backed by a local directory tree, ``/object_store`` unless ``root`` is given. A write lands whole or not at all, and keeps any content type but the default. The root-level ``.gitignore``, - ``.agentenv-meta`` and ``.agentenv-tmp`` keys are reserved for the store's own files.""" + ``.agentenv-meta`` and ``.agentenv-tmp`` keys are reserved for the store's own files. + + Its HTTPS transfer grants are served by a server in this process, started by the first grant: it listens on + ``grant_bind_host`` (by default loopback, or the Docker bridge gateway on Linux) and grant URLs name + ``grant_advertise_host`` (by default ``host.docker.internal``), which must be a local name or a loopback or + private address. A grant works until it expires or this process exits.""" - def __init__(self, root: str) -> None: - self._root = Path(root) + def __init__( + self, root: str | None = None, *, grant_bind_host: str | None = None, grant_advertise_host: str | None = None + ) -> None: + if grant_bind_host is not None: + try: + ipaddress.ip_address(grant_bind_host) + except ValueError: + raise ValueError(f"grant_bind_host must be an IP address, got {grant_bind_host!r}") from None + if grant_advertise_host is not None: + check_local_host(grant_advertise_host) + self._root = Path(root) if root is not None else state_root() / "object_store" + self._grant_bind_host = grant_bind_host + self._grant_advertise_host = grant_advertise_host @property def root(self) -> Path: @@ -126,6 +155,41 @@ def get_object_key(self, object_url: str) -> str: raise ValueError(f"{object_url!r} is not an object in {root}.") return path.relative_to(root).as_posix() + def issue_read_grant(self, object_url: str, *, expires_in: int = 3600) -> HttpGetGrant: + url, expires_at = self._grant("get", expires_in, key=self._grant_key(object_url)) + return HttpGetGrant(kind="http-get", url=url, expires_at=expires_at) + + def issue_write_grant( + self, object_url: str, *, media_type: str, max_bytes: int, expires_in: int = 3600 + ) -> HttpPutGrant: + url, expires_at = self._grant( + "put", expires_in, key=self._grant_key(object_url), max_bytes=max_bytes, content_type=media_type + ) + return HttpPutGrant(kind="http-put", url=url, expires_at=expires_at, headers={"Content-Type": media_type}) + + def issue_upload_policy(self, prefix_url: str, *, max_object_bytes: int, expires_in: int) -> UploadPolicy: + key = self.get_object_key(prefix_url) + url, expires_at = self._grant( + "post", expires_in, prefix="" if key == "." else key + "/", max_bytes=max_object_bytes + ) + return UploadPolicy( + write=HttpPostPolicyGrant(kind="http-post-policy", url=url, fields={}, path_field="key", file_field="file"), + expires_at=expires_at, + ) + + def _grant(self, op: Op, expires_in: int, **scope) -> tuple[str, datetime]: + if expires_in <= 0: + raise ValueError(f"expires_in must be positive, got {expires_in}") + expires = int(time.time()) + expires_in + server = grant_server(self._grant_bind_host, self._grant_advertise_host) + return server.issue(self, op, expires, **scope), datetime.fromtimestamp(expires, timezone.utc) + + def _grant_key(self, object_url: str) -> str: + """The key of the one object at ``object_url``, refusing the store's own files.""" + key = self.get_object_key(object_url) + self._resolve(key) + return key + def _resolve(self, key: str) -> Path: root = self._root.resolve() resolved = (self._root / key).resolve() diff --git a/tst/unit/store/local_grants_server_test.py b/tst/unit/store/local_grants_server_test.py new file mode 100644 index 00000000..fd53b1db --- /dev/null +++ b/tst/unit/store/local_grants_server_test.py @@ -0,0 +1,270 @@ +"""The local store's grants against its real grant server, over HTTPS on loopback: the conformance +grant cases, and the bounds and refusals S3 would apply to the same requests.""" + +import subprocess +import time + +import httpx +import pytest +from pytest_socket import enable_socket + +from agent_env.config.paths import state_root +from agent_env.store import LocalFilesystemObjectStore +from agent_env.store.object_store.local_grants import server as server_module +from agent_env.store.object_store.local_grants.server import default_bind_host, grant_server +from agent_env.store.object_store.local_grants.tls import local_ca +from tst.store import object_conformance + + +@pytest.fixture(autouse=True) +def _disable_network(): + """Override the suite-wide socket block (tst/unit/conftest.py) for this module: loopback is the point here.""" + enable_socket() + yield + + +@pytest.fixture +def store(tmp_path, monkeypatch): + """A local store granting on loopback under this test's state root, trusted as a container would trust it.""" + monkeypatch.setattr(server_module, "_servers", {}) + monkeypatch.setenv("SSL_CERT_FILE", str(local_ca().bundle_path)) + yield LocalFilesystemObjectStore(str(tmp_path / "objects"), grant_bind_host="127.0.0.1", grant_advertise_host="localhost") + for server in server_module._servers.values(): + server.close() + + +def _policy(store, prefix="ns/", max_object_bytes=16): + return store.issue_upload_policy(store.object_url(prefix), max_object_bytes=max_object_bytes, expires_in=600).write + + +def _post(write, fields, data=b"data", *, file_first=False): + files = {write.file_field: ("object", data, "application/octet-stream")} + if file_first: # httpx sends data fields before files, so build the body by hand + request = httpx.Request("POST", write.url, files=files) + body = request.read() + boundary = request.headers["content-type"].split("boundary=")[1] + extra = b"".join( + f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"\r\n\r\n{v}\r\n'.encode() for k, v in fields.items() + ) + body = body.replace(f"--{boundary}--".encode(), extra + f"--{boundary}--".encode()) + return httpx.post(write.url, content=body, headers={"Content-Type": request.headers["content-type"]}) + return httpx.post(write.url, data=fields, files=files) + + +@pytest.mark.parametrize("case", object_conformance.GRANT_CASES, ids=lambda c: c.__name__) +def test_grant_conformance(case, store): + case(store, "") + + +def test_a_read_returns_the_exact_bytes_and_type(store): + url = store.put("a/page.html", b"

hi

", content_type="application/json") + response = httpx.get(store.issue_read_grant(url).url, headers={"Accept-Encoding": "gzip"}) + assert response.status_code == 200 + assert response.content == b"

hi

" + assert response.headers["content-type"] == "application/json" + assert response.headers["content-length"] == "9" + assert "content-encoding" not in response.headers + + +def test_a_read_of_a_missing_object_is_404(store): + assert httpx.get(store.issue_read_grant(store.object_url("a/none.bin")).url).status_code == 404 + + +def test_a_head_reports_the_size(store): + url = store.put("a/b.bin", b"12345") + response = httpx.head(store.issue_read_grant(url).url) + assert response.status_code == 200 + assert response.headers["content-length"] == "5" + + +def test_a_grant_allows_only_its_own_method(store): + url = store.put("a/b.bin", b"v") + read = store.issue_read_grant(url).url + assert httpx.put(read, content=b"x").status_code == 403 + assert httpx.post(read, content=b"x").status_code == 403 + assert httpx.delete(read).status_code == 405 + assert store.get(url) == b"v" + + +def test_a_tampered_or_expired_grant_is_403(store): + url = store.put("a/b.bin", b"v") + read = store.issue_read_grant(url).url + assert httpx.get(read[:-3] + "AAA").status_code == 403 + assert httpx.get(read.rsplit("/", 1)[0] + "/not-a-token").status_code == 403 + server = grant_server("127.0.0.1", "localhost") + expired = server.issue(store, "get", int(time.time()) - 1, key="a/b.bin") + assert httpx.get(expired).status_code == 403 + + +def test_a_grant_from_a_closed_server_stops_working(store): + url = store.put("a/b.bin", b"v") + read = store.issue_read_grant(url).url + grant_server("127.0.0.1", "localhost").close() + with pytest.raises(httpx.ConnectError): + httpx.get(read) + assert httpx.get(store.issue_read_grant(url).url).content == b"v" # the next grant starts it again + + +def test_a_write_must_declare_the_signed_type(store): + url = store.object_url("a/w.json") + grant = store.issue_write_grant(url, media_type="application/json", max_bytes=64) + assert httpx.put(grant.url, content=b"{}", headers={"Content-Type": "text/plain"}).status_code == 403 + assert httpx.put(grant.url, content=b"{}").status_code == 403 + assert not store.exists("a/w.json") + + +def test_a_write_larger_than_the_grant_is_refused(store): + url = store.object_url("a/w.bin") + grant = store.issue_write_grant(url, media_type="application/octet-stream", max_bytes=4) + assert httpx.put(grant.url, content=b"12345", headers=grant.headers).status_code == 400 + + def chunks(): # no Content-Length: the bound is enforced as the body streams + yield b"123" + yield b"45" + + assert httpx.put(grant.url, content=chunks(), headers=grant.headers).status_code == 400 + assert not store.exists("a/w.bin") + assert not any((store.root / ".agentenv-tmp").iterdir()) + + +def test_a_write_grant_may_be_retried_and_overwrites(store): + url = store.object_url("a/w.bin") + grant = store.issue_write_grant(url, media_type="application/octet-stream", max_bytes=8) + assert httpx.put(grant.url, content=b"first", headers=grant.headers).status_code == 200 + assert httpx.put(grant.url, content=b"second", headers=grant.headers).status_code == 200 + assert store.get(url) == b"second" + + +def test_an_upload_keeps_its_content_type(store): + write = _policy(store) + assert _post(write, {"key": "ns/x.txt", "Content-Type": "text/plain"}, b"hello").status_code == 204 + assert store.read("ns/x.txt") == b"hello" + assert store.get_object_metadata("ns/x.txt").content_type == "text/plain" + + +def test_an_upload_may_nest_and_overwrite(store): + write = _policy(store) + assert _post(write, {"key": "ns/a/b/c"}, b"one").status_code == 204 + assert _post(write, {"key": "ns/a/b/c"}, b"two").status_code == 204 + assert store.read("ns/a/b/c") == b"two" + + +@pytest.mark.parametrize("key", ["ns/../x", "ns//x", "ns/./x", "/ns/x", "ns\\x", "ns-sibling/x", "other/x", "ns"]) +def test_an_upload_key_outside_the_prefix_is_403(store, key): + assert _post(_policy(store), {"key": key}).status_code == 403 + assert store.list("") == [] + + +def test_an_upload_at_the_root_cannot_touch_the_stores_own_files(store): + write = _policy(store, prefix="") + assert _post(write, {"key": ".agentenv-meta/x"}).status_code == 403 + assert _post(write, {"key": "top.bin"}).status_code == 204 + assert store.read("top.bin") == b"data" + + +def test_an_upload_larger_than_the_grant_is_refused(store): + write = _policy(store, max_object_bytes=4) + assert _post(write, {"key": "ns/big"}, b"12345").status_code == 400 + assert not store.exists("ns/big") + assert not any((store.root / ".agentenv-tmp").iterdir()) + + +def test_an_upload_must_name_its_key_before_its_file(store): + write = _policy(store) + assert _post(write, {"key": "ns/late"}, file_first=True).status_code == 400 + assert not store.exists("ns/late") + + +def test_fields_after_the_file_are_ignored(store): + write = _policy(store) + request = httpx.Request("POST", write.url, data={"key": "ns/second"}, files={"file": ("f", b"v", "text/plain")}) + boundary = request.headers["content-type"].split("boundary=")[1] + body = request.read().replace( + f"--{boundary}--".encode(), + f'--{boundary}\r\nContent-Disposition: form-data; name="key"\r\n\r\nns/../escape\r\n--{boundary}--'.encode(), + ) + assert httpx.post(write.url, content=body, headers={"Content-Type": request.headers["content-type"]}).status_code == 204 + assert store.read("ns/second") == b"v" + + +@pytest.mark.parametrize("body, content_type", [ + (b"key=ns/x", "application/x-www-form-urlencoded"), + (b"--b\r\nContent-Disposition: form-data; name=\"key\"\r\n\r\nns/x\r\n--b--\r\n", "multipart/form-data; boundary=b"), +]) +def test_an_upload_that_is_not_a_form_with_a_file_is_400(store, body, content_type): + write = _policy(store) + assert httpx.post(write.url, content=body, headers={"Content-Type": content_type}).status_code == 400 + + +def test_grants_refuse_the_stores_own_files(store): + with pytest.raises(ValueError, match="reserved"): + store.issue_read_grant(store.object_url(".agentenv-meta/a")) + + +def test_a_grant_expires_when_asked(store): + url = store.put("a/b.bin", b"v") + grant = store.issue_read_grant(url, expires_in=120) + assert 118 <= (grant.expires_at.timestamp() - time.time()) <= 121 + with pytest.raises(ValueError): + store.issue_read_grant(url, expires_in=0) + + +def test_two_stores_share_one_server_and_keep_their_own_objects(store, tmp_path): + other = LocalFilesystemObjectStore(str(tmp_path / "other"), grant_bind_host="127.0.0.1", grant_advertise_host="localhost") + a, b = store.put("k", b"mine"), other.put("k", b"theirs") + ga, gb = store.issue_read_grant(a).url, other.issue_read_grant(b).url + assert ga.split("/v1/")[0] == gb.split("/v1/")[0] + assert (httpx.get(ga).content, httpx.get(gb).content) == (b"mine", b"theirs") + + +def test_grant_urls_name_the_advertised_host(tmp_path, monkeypatch): + monkeypatch.setattr(server_module, "_servers", {}) + store = LocalFilesystemObjectStore(str(tmp_path), grant_bind_host="127.0.0.1") + try: + assert store.issue_read_grant(store.put("k", b"v")).url.startswith("https://host.docker.internal:") + finally: + for server in server_module._servers.values(): + server.close() + + +def test_settings_are_checked_when_the_store_is_made(tmp_path): + with pytest.raises(ValueError, match="IP address"): + LocalFilesystemObjectStore(str(tmp_path), grant_bind_host="localhost") + with pytest.raises(ValueError, match="not a local name"): + LocalFilesystemObjectStore(str(tmp_path), grant_advertise_host="example.com") + with pytest.raises(ValueError, match="private address"): + LocalFilesystemObjectStore(str(tmp_path), grant_advertise_host="8.8.8.8") + + +def test_the_root_defaults_to_the_state_root(): + assert LocalFilesystemObjectStore().root == state_root() / "object_store" + + +class TestDefaultBindHost: + @pytest.fixture(autouse=True) + def _fresh(self): + default_bind_host.cache_clear() + yield + default_bind_host.cache_clear() + + def test_loopback_off_linux(self, monkeypatch): + monkeypatch.setattr(server_module.platform, "system", lambda: "Darwin") + assert default_bind_host() == "127.0.0.1" + + def test_the_bridge_gateway_on_linux(self, monkeypatch): + monkeypatch.setattr(server_module.platform, "system", lambda: "Linux") + monkeypatch.setattr( + server_module.subprocess, "run", + lambda *a, **k: subprocess.CompletedProcess(a, 0, stdout="172.17.0.1 \n", stderr=""), + ) + assert default_bind_host() == "172.17.0.1" + + def test_loopback_with_a_warning_when_docker_cannot_say(self, monkeypatch, caplog): + monkeypatch.setattr(server_module.platform, "system", lambda: "Linux") + + def fail(*a, **k): + raise FileNotFoundError("docker") + + monkeypatch.setattr(server_module.subprocess, "run", fail) + assert default_bind_host() == "127.0.0.1" + assert "grant_bind_host" in caplog.text diff --git a/tst/unit/store/local_grants_tls_test.py b/tst/unit/store/local_grants_tls_test.py new file mode 100644 index 00000000..b1c68a04 --- /dev/null +++ b/tst/unit/store/local_grants_tls_test.py @@ -0,0 +1,107 @@ +"""The local transfer CA: created once per user, trusted only through its own files, and able to +vouch for local names and private addresses only. Handshakes run over memory BIOs, not sockets.""" + +import os +import ssl +import stat +from datetime import timedelta +from pathlib import Path + +import certifi +import pytest +from cryptography import x509 + +from agent_env.config.paths import state_root +from agent_env.store.object_store.local_grants import tls +from agent_env.store.object_store.local_grants.tls import check_local_host, local_ca, server_context + + +def _handshake(server: ssl.SSLContext, *, cafile: Path, hostname: str) -> None: + client = ssl.create_default_context(cafile=str(cafile)) + client.verify_flags |= ssl.VERIFY_X509_STRICT # the default from Python 3.13 + to_server, to_client = ssl.MemoryBIO(), ssl.MemoryBIO() + sides = [ + client.wrap_bio(to_client, to_server, server_hostname=hostname), + server.wrap_bio(to_server, to_client, server_side=True), + ] + done = set() + for _ in range(10): + for side in sides: + if id(side) not in done: + try: + side.do_handshake() + done.add(id(side)) + except ssl.SSLWantReadError: + pass + if len(done) == 2: + return + raise AssertionError("the handshake did not finish") + + +@pytest.mark.parametrize("hostname", ["host.docker.internal", "localhost", "127.0.0.1"]) +def test_a_server_certificate_verifies_for_the_names_it_was_made_for(hostname): + ca = local_ca() + context = server_context(ca, ["localhost", "host.docker.internal", "127.0.0.1"]) + _handshake(context, cafile=ca.cert_path, hostname=hostname) + _handshake(context, cafile=ca.bundle_path, hostname=hostname) + + +def test_the_ca_cannot_vouch_for_a_public_name(): + ca = local_ca() + with pytest.raises(ssl.SSLCertVerificationError): + _handshake(server_context(ca, ["example.com"]), cafile=ca.cert_path, hostname="example.com") + + +def test_the_ca_cannot_vouch_for_a_public_address(): + ca = local_ca() + with pytest.raises(ssl.SSLCertVerificationError): + _handshake(server_context(ca, ["8.8.8.8"]), cafile=ca.cert_path, hostname="8.8.8.8") + + +def test_the_bundle_is_the_public_roots_plus_the_ca(): + ca = local_ca() + bundle = ca.bundle_path.read_bytes() + assert bundle.startswith(Path(certifi.where()).read_bytes().rstrip()) + assert bundle.endswith(ca.cert_path.read_bytes()) + assert x509.load_pem_x509_certificate(ca.cert_path.read_bytes()) == ca.cert + + +def test_the_ca_key_is_private_to_the_user(): + local_ca() + tls_dir = state_root() / "tls" + assert stat.S_IMODE(os.stat(tls_dir).st_mode) == 0o700 + assert stat.S_IMODE(os.stat(tls_dir / "ca.key.pem").st_mode) == 0o600 + + +def test_the_ca_is_created_once_and_reused(): + assert local_ca().cert == local_ca().cert + + +def test_a_ca_near_expiry_is_replaced(monkeypatch): + with monkeypatch.context() as m: + m.setattr(tls, "_CA_LIFETIME", timedelta(days=10)) + first = local_ca() + second = local_ca() + assert second.cert != first.cert + assert x509.load_pem_x509_certificate(second.cert_path.read_bytes()) == second.cert + assert local_ca().cert == second.cert + + +def test_an_unreadable_ca_is_replaced(): + first = local_ca() + (state_root() / "tls" / "ca.key.pem").write_text("not a key") + assert local_ca().cert != first.cert + + +@pytest.mark.parametrize("host", [ + "localhost", "host.docker.internal", "a.localhost", "127.0.0.1", "172.17.0.1", "192.168.5.2", "10.0.0.8", + "::1", "fd00::1", +]) +def test_local_hosts_are_accepted(host): + check_local_host(host) + + +@pytest.mark.parametrize("host", ["example.com", "evil-localhost", "docker.internal", "8.8.8.8", "2001:db8::1"]) +def test_other_hosts_are_refused(host): + with pytest.raises(ValueError): + check_local_host(host) diff --git a/tst/unit/store/local_grants_tokens_test.py b/tst/unit/store/local_grants_tokens_test.py new file mode 100644 index 00000000..9ed7d427 --- /dev/null +++ b/tst/unit/store/local_grants_tokens_test.py @@ -0,0 +1,65 @@ +"""Grant tokens: what one signs, it verifies, until the expiry; nothing else verifies.""" + +import pytest + +from agent_env.store.object_store.local_grants.tokens import GrantClaims, GrantSigner, InvalidGrantError + +NOW = 1_800_000_000 + + +def _claims(**overrides) -> GrantClaims: + fields = {"op": "put", "store": "s1", "expires": NOW + 60, "key": "a/b.json", "max_bytes": 64, + "content_type": "application/json"} + return GrantClaims(**{**fields, **overrides}) + + +def test_a_signed_token_verifies_to_its_claims(): + signer = GrantSigner() + claims = _claims() + assert signer.verify(signer.sign(claims), now=NOW) == claims + + +def test_a_post_token_carries_its_prefix(): + signer = GrantSigner() + claims = GrantClaims(op="post", store="s1", expires=NOW + 60, prefix="ns/", max_bytes=8) + assert signer.verify(signer.sign(claims), now=NOW).prefix == "ns/" + + +def test_an_expired_token_is_refused(): + signer = GrantSigner() + token = signer.sign(_claims(expires=NOW)) + with pytest.raises(InvalidGrantError, match="expired"): + signer.verify(token, now=NOW) + + +def test_another_processs_token_is_refused(): + with pytest.raises(InvalidGrantError, match="signature"): + GrantSigner().verify(GrantSigner().sign(_claims()), now=NOW) + + +@pytest.mark.parametrize("tamper", [ + lambda t: t.replace("v1.", "v2.", 1), + lambda t: t.rsplit(".", 1)[0], + lambda t: t[:-2] + ("AA" if not t.endswith("AA") else "BB"), + lambda t: "v1." + GrantSigner().sign(_claims(key="other")).split(".")[1] + "." + t.rsplit(".", 1)[1], + lambda t: t + "!", +]) +def test_a_tampered_token_is_refused(tamper): + signer = GrantSigner() + with pytest.raises(InvalidGrantError): + signer.verify(tamper(signer.sign(_claims())), now=NOW) + + +@pytest.mark.parametrize("fields", [ + {"op": "delete"}, + {"op": "post", "key": None, "prefix": "ns"}, + {"op": "post", "prefix": "ns/"}, + {"op": "get", "key": None, "prefix": "ns/"}, + {"max_bytes": None}, + {"max_bytes": -1}, + {"content_type": None}, + {"op": "get", "max_bytes": None}, +]) +def test_claims_that_make_no_sense_are_refused(fields): + with pytest.raises(ValueError): + _claims(**fields) diff --git a/tst/unit/store/test_config_object_store.py b/tst/unit/store/test_config_object_store.py index c3d2db6d..3a5a996f 100644 --- a/tst/unit/store/test_config_object_store.py +++ b/tst/unit/store/test_config_object_store.py @@ -83,6 +83,20 @@ def test_config_toml_selects_custom_impl(monkeypatch, tmp_path): assert isinstance(Config().get_object_store(), FakeObjectStore) +def test_config_toml_sets_the_local_stores_grant_hosts(monkeypatch, tmp_path): + monkeypatch.delenv("AGENT_ENV_OBJECT_STORE", raising=False) + monkeypatch.delenv("AGENT_ENV_CONFIG", raising=False) + monkeypatch.chdir(tmp_path) + _write_config(tmp_path, ( + '[stores.object]\nimpl = "agent_env.store.object_store:LocalFilesystemObjectStore"\n' + 'config = { grant_bind_host = "172.18.0.1", grant_advertise_host = "172.18.0.1" }\n' + )) + store = Config().get_object_store() + assert isinstance(store, LocalFilesystemObjectStore) + assert store.root == state_root() / "object_store" + assert (store._grant_bind_host, store._grant_advertise_host) == ("172.18.0.1", "172.18.0.1") + + def test_env_override_beats_config_toml(monkeypatch, tmp_path): monkeypatch.delenv("AGENT_ENV_CONFIG", raising=False) monkeypatch.chdir(tmp_path) diff --git a/tst/unit/store/test_object_transfer_grants.py b/tst/unit/store/test_object_transfer_grants.py index 2af18057..9027dd56 100644 --- a/tst/unit/store/test_object_transfer_grants.py +++ b/tst/unit/store/test_object_transfer_grants.py @@ -281,16 +281,9 @@ def test_botocore_still_exposes_the_signing_credentials() -> None: assert S3ObjectStore(session, BUCKET)._signing_credentials() == (None, True) -def test_local_store_cannot_issue_remote_transfer_grants(tmp_path) -> None: +def test_local_store_does_not_offer_transfer_grants(tmp_path) -> None: + """It can issue them (local_grants_server_test.py), but agents on the local store don't use them.""" store = LocalFilesystemObjectStore(str(tmp_path)) assert S3ObjectStore.supports_transfer_grants assert not store.supports_transfer_grants - with pytest.raises(NotImplementedError, match="cannot issue"): - store.issue_read_grant(store.object_url("input")) - with pytest.raises(NotImplementedError, match="cannot issue"): - store.issue_write_grant( - store.object_url("output"), media_type="application/json", max_bytes=10 - ) - with pytest.raises(NotImplementedError, match="cannot issue"): - store.issue_upload_policy(store.object_url("changelog"), max_object_bytes=10, expires_in=60) diff --git a/uv.lock b/uv.lock index 0cb48569..5e46e292 100644 --- a/uv.lock +++ b/uv.lock @@ -49,7 +49,9 @@ dependencies = [ { name = "a2a-sdk" }, { name = "agentenv-protocol" }, { name = "boto3" }, + { name = "certifi" }, { name = "click" }, + { name = "cryptography" }, { name = "e2b" }, { name = "httpx" }, { name = "litellm" }, @@ -58,7 +60,10 @@ dependencies = [ { name = "packaging" }, { name = "pydantic" }, { name = "pymongo" }, + { name = "python-multipart" }, { name = "pyyaml" }, + { name = "starlette" }, + { name = "uvicorn" }, ] [package.optional-dependencies] @@ -98,7 +103,9 @@ requires-dist = [ { name = "agentenv-framework", extras = ["explorer", "gcp"], marker = "extra == 'dev'" }, { name = "agentenv-protocol", editable = "packages/agentenv-protocol" }, { name = "boto3", specifier = ">=1.35.0" }, + { name = "certifi", specifier = ">=2024.2.2" }, { name = "click", specifier = ">=8.0.0" }, + { name = "cryptography", specifier = ">=42" }, { name = "e2b", specifier = "~=2.46.4" }, { name = "fastapi", marker = "extra == 'explorer'", specifier = ">=0.110" }, { name = "google-api-core", marker = "extra == 'gcp'", specifier = ">=2.24" }, @@ -121,8 +128,11 @@ requires-dist = [ { name = "pytest-socket", marker = "extra == 'dev'", specifier = ">=0.7.0" }, { name = "pytest-timeout", marker = "extra == 'dev'", specifier = ">=2.3.0" }, { name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.0.0" }, + { name = "python-multipart", specifier = ">=0.0.13" }, { name = "pyyaml", specifier = ">=6.0" }, { name = "requests", marker = "extra == 'gcp'", specifier = ">=2.31" }, + { name = "starlette", specifier = ">=0.40" }, + { name = "uvicorn", specifier = ">=0.47" }, { name = "uvicorn", marker = "extra == 'explorer'", specifier = ">=0.29" }, ] provides-extras = ["explorer", "gcp", "dev"] From a9af97e096f5ba7fff7d1f6aa07adec88e81bb56 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 07:21:45 -1000 Subject: [PATCH 03/14] fix(store): keep the local CA's public trust files in their own readable directory ca.pem and ca-bundle.pem now live in /tls/trust/, mode 0755, apart from the CA key, so the directory can be copied into a container whole and read by whatever user it runs as. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_grants/tls.py | 16 +++++++++++----- tst/unit/store/local_grants_tls_test.py | 7 +++++-- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/agent_env/store/object_store/local_grants/tls.py b/src/agent_env/store/object_store/local_grants/tls.py index d595db5d..13b5c764 100644 --- a/src/agent_env/store/object_store/local_grants/tls.py +++ b/src/agent_env/store/object_store/local_grants/tls.py @@ -4,7 +4,8 @@ local names and private addresses, so even a leaked key can only vouch for those, and it is never installed in the host's trust stores: only the containers agent-env starts trust it, through ``ca-bundle.pem`` (the public roots plus this CA, for ``SSL_CERT_FILE``, which replaces -the defaults) or ``ca.pem`` (the CA alone, for trust stores that add to the defaults). +the defaults) or ``ca.pem`` (the CA alone, for trust stores that add to the defaults). Those two +are all that ``tls/trust/`` holds, so it can be copied into a container whole. """ from __future__ import annotations @@ -41,6 +42,7 @@ _CLOCK_SKEW = timedelta(minutes=5) _CA_FILE = "ca.key.pem" # the CA's key and certificate, private to the user +_TRUST_DIR = "trust" # the public files, readable by any user a container runs as _CERT_FILE = "ca.pem" _BUNDLE_FILE = "ca-bundle.pem" _PEM_CERT = b"-----BEGIN CERTIFICATE-----" @@ -48,10 +50,11 @@ @dataclass(frozen=True) class LocalCA: - """The CA's key and certificate, and the trust files derived from it.""" + """The CA's key and certificate, and the trust files derived from it in ``trust_dir``.""" key: ec.EllipticCurvePrivateKey cert: x509.Certificate + trust_dir: Path cert_path: Path bundle_path: Path @@ -96,10 +99,13 @@ def local_ca() -> LocalCA: staged.unlink(missing_ok=True) loaded = _load(ca_file) or (key, cert) key, cert = loaded + trust_dir = tls_dir / _TRUST_DIR + trust_dir.mkdir(exist_ok=True) + trust_dir.chmod(0o755) cert_pem = cert.public_bytes(serialization.Encoding.PEM) - cert_path = _write_if_changed(tls_dir / _CERT_FILE, cert_pem) - bundle_path = _write_if_changed(tls_dir / _BUNDLE_FILE, Path(certifi.where()).read_bytes().rstrip() + b"\n" + cert_pem) - return LocalCA(key=key, cert=cert, cert_path=cert_path, bundle_path=bundle_path) + cert_path = _write_if_changed(trust_dir / _CERT_FILE, cert_pem) + bundle_path = _write_if_changed(trust_dir / _BUNDLE_FILE, Path(certifi.where()).read_bytes().rstrip() + b"\n" + cert_pem) + return LocalCA(key=key, cert=cert, trust_dir=trust_dir, cert_path=cert_path, bundle_path=bundle_path) def server_context(ca: LocalCA, hosts: Iterable[str]) -> ssl.SSLContext: diff --git a/tst/unit/store/local_grants_tls_test.py b/tst/unit/store/local_grants_tls_test.py index b1c68a04..d209f443 100644 --- a/tst/unit/store/local_grants_tls_test.py +++ b/tst/unit/store/local_grants_tls_test.py @@ -66,11 +66,14 @@ def test_the_bundle_is_the_public_roots_plus_the_ca(): assert x509.load_pem_x509_certificate(ca.cert_path.read_bytes()) == ca.cert -def test_the_ca_key_is_private_to_the_user(): - local_ca() +def test_the_ca_key_is_private_to_the_user_and_the_trust_files_are_not(): + ca = local_ca() tls_dir = state_root() / "tls" assert stat.S_IMODE(os.stat(tls_dir).st_mode) == 0o700 assert stat.S_IMODE(os.stat(tls_dir / "ca.key.pem").st_mode) == 0o600 + assert stat.S_IMODE(os.stat(ca.trust_dir).st_mode) == 0o755 + assert sorted(p.name for p in ca.trust_dir.iterdir()) == ["ca-bundle.pem", "ca.pem"] + assert {stat.S_IMODE(os.stat(p).st_mode) for p in ca.trust_dir.iterdir()} == {0o644} def test_the_ca_is_created_once_and_reused(): From e91eea98855b490b86cb3cac094b5bf4f7e32252 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 07:51:22 -1000 Subject: [PATCH 04/14] fix(store): a recorded content type names the file it describes, and a write-once put never exposes an empty object A content type is recorded with the identity of the file it was written for (inode, size, mtime, which a rename keeps), so a type recorded for another write - a racing overwrite that renamed first, or an old record an overwrite could not remove - reads back as unknown, never as the wrong type. On a filesystem without hard links a write-once put now checks for the key and renames into place, instead of reserving the key with an empty file that readers could see and a crash could leave. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 33 ++++++++++++------- tst/unit/store/local_object_store_test.py | 32 ++++++++++++++++++ 2 files changed, 54 insertions(+), 11 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index 5db87829..d6e10486 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -23,7 +23,7 @@ _META_DIR = ".agentenv-meta" # each object's content type, at the object's own key _STAGING_DIR = ".agentenv-tmp" # files being written, inside the root so a rename into place is atomic _RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) -# A filesystem without hard links: a no-overwrite write reserves its key by exclusive create instead. +# A filesystem without hard links: a no-overwrite write checks, then renames, so two racing writers can both land. _NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) @@ -92,7 +92,7 @@ def get_object_metadata(self, key: str) -> ObjectMetadata | None: return None st = path.stat() return ObjectMetadata( - content_type=self._read_content_type(path), + content_type=self._read_content_type(path, st), size=st.st_size, last_modified=datetime.fromtimestamp(st.st_mtime, tz=timezone.utc), ) @@ -166,7 +166,11 @@ def _staged(self) -> Iterator[Path]: def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrite: bool) -> None: """Put a staged file in place at ``path`` and record its content type. An overwrite renames over - whatever is there; otherwise a hard link claims ``path`` only if no other writer got there first.""" + whatever is there; otherwise a hard link claims ``path`` only if no other writer got there first. + + The type is recorded with the identity of the file it describes, so one recorded for another + write (a racing overwrite, or one whose type could not be recorded) reads back as unknown.""" + identity = _identity(staged.stat()) path.parent.mkdir(parents=True, exist_ok=True) if allow_overwrite: os.replace(staged, path) @@ -178,9 +182,7 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit except OSError as e: if e.errno not in _NO_HARD_LINKS: raise - try: - os.close(os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)) - except FileExistsError: + if path.exists(): raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None os.replace(staged, path) # The default type is what a writer that named none gets; it reads back as unknown, so readers @@ -192,15 +194,19 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit return meta.parent.mkdir(parents=True, exist_ok=True) with self._staged() as staged_meta: - staged_meta.write_text(json.dumps({"content_type": content_type})) + staged_meta.write_text(json.dumps({"content_type": content_type, "object": identity})) os.replace(staged_meta, meta) - except OSError as e: # the object is in place; only its type goes unrecorded + except OSError as e: # the object is in place and its type reads back as unknown logger.warning("Could not record the content type of %s: %s", path, e) - def _read_content_type(self, path: Path) -> str | None: - """The content type recorded for the object at ``path``; None for one written before types were kept.""" + def _read_content_type(self, path: Path, st: os.stat_result | None = None) -> str | None: + """The content type recorded for the object at ``path``, if it was recorded for this very file; None + otherwise, as for one written before types were kept.""" try: - return json.loads(self._meta_path(path).read_text()).get("content_type") + recorded = json.loads(self._meta_path(path).read_text()) + if recorded.get("object") != _identity(st or path.stat()): + return None + return recorded.get("content_type") except (OSError, ValueError, AttributeError): return None @@ -217,6 +223,11 @@ def _from_url(object_url: str) -> str: return object_url[len("file://"):] if object_url.startswith("file://") else object_url +def _identity(st: os.stat_result) -> list[int]: + """What tells one file apart from another written to the same key: a rename keeps all three.""" + return [st.st_ino, st.st_size, st.st_mtime_ns] + + @contextlib.contextmanager def _missing_as_not_found(object_url: str) -> Iterator[None]: try: diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index 3ac8e048..e8295c63 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -176,6 +176,38 @@ def no_links(src, dst): assert store.read("nolink/x") == b"first" +def test_without_hard_links_a_failed_write_once_put_leaves_no_object(store, monkeypatch): + def no_links(src, dst): + raise OSError(errno.EPERM, "hard links not supported") + + def failing_replace(src, dst): + raise OSError(errno.EIO, "disk went away") + + monkeypatch.setattr(os, "link", no_links) + monkeypatch.setattr(os, "replace", failing_replace) + with pytest.raises(OSError, match="disk went away"): + store.put("nolink/y", b"never") + assert not store.exists("nolink/y") + + +def test_a_type_recorded_for_another_write_reads_back_as_unknown(store, tmp_path): + """Two overwrites of one key race: the one that renamed first records its type last. Its type must not + be read as the other's.""" + store.put("race/x", b"first", content_type="text/plain") + first_meta = (tmp_path / ".agentenv-meta" / "race" / "x").read_text() + store.put("race/x", b"second!", content_type="application/json", allow_overwrite=True) + (tmp_path / ".agentenv-meta" / "race" / "x").write_text(first_meta) + metadata = store.get_object_metadata("race/x") + assert (metadata.content_type, metadata.size) == (None, 7) + + +def test_an_overwrite_that_names_no_type_never_inherits_the_old_one(store, tmp_path, monkeypatch): + store.put("keep/x", b"typed", content_type="text/plain") + monkeypatch.setattr(Path, "unlink", lambda self, missing_ok=False: None) # the old record survives + store.put("keep/x", b"untyped!", allow_overwrite=True) + assert store.get_object_metadata("keep/x").content_type is None + + @pytest.mark.parametrize("key", [".agentenv-meta/x", ".agentenv-tmp/x", "a/../.agentenv-meta"]) def test_the_stores_own_directories_are_reserved(store, key): with pytest.raises(ValueError, match="reserved"): From 38846c9021244eb855f8dba39d237686e86678e9 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 07:53:41 -1000 Subject: [PATCH 05/14] fix(store): a process keeps one local CA, and grant uploads bound their part headers A process now keeps the local CA it first loads, and each CA's trust files are written once into a directory of their own under tls/trust/, so the grant server and the containers a process starts agree on the CA even when another process renews the one on disk. A server certificate lasts as long as its CA rather than 90 days, since a server keeps one for its whole life. A multipart upload's part headers are bounded to 16 KiB, as its fields and file already are. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_grants/server.py | 10 +++++ .../store/object_store/local_grants/tls.py | 24 ++++++++---- tst/unit/store/local_grants_server_test.py | 11 ++++++ tst/unit/store/local_grants_tls_test.py | 37 ++++++++++++++++++- 4 files changed, 73 insertions(+), 9 deletions(-) diff --git a/src/agent_env/store/object_store/local_grants/server.py b/src/agent_env/store/object_store/local_grants/server.py index cf2449ce..759966fd 100644 --- a/src/agent_env/store/object_store/local_grants/server.py +++ b/src/agent_env/store/object_store/local_grants/server.py @@ -53,6 +53,7 @@ _UPLOAD_FILE_FIELD = "file" _MAX_FORM_FIELDS = 32 _MAX_FORM_FIELD_BYTES = 64 * 1024 +_MAX_PART_HEADER_BYTES = 16 * 1024 _servers: dict[tuple[str, str], GrantServer] = {} _servers_lock = threading.Lock() @@ -289,6 +290,7 @@ def __init__(self, store: LocalFilesystemObjectStore, claims: GrantClaims, stack self._stack = stack 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 @@ -319,15 +321,23 @@ def commit(self) -> None: 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 _Rejected(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"" diff --git a/src/agent_env/store/object_store/local_grants/tls.py b/src/agent_env/store/object_store/local_grants/tls.py index 13b5c764..d3ae6d90 100644 --- a/src/agent_env/store/object_store/local_grants/tls.py +++ b/src/agent_env/store/object_store/local_grants/tls.py @@ -4,12 +4,14 @@ local names and private addresses, so even a leaked key can only vouch for those, and it is never installed in the host's trust stores: only the containers agent-env starts trust it, through ``ca-bundle.pem`` (the public roots plus this CA, for ``SSL_CERT_FILE``, which replaces -the defaults) or ``ca.pem`` (the CA alone, for trust stores that add to the defaults). Those two -are all that ``tls/trust/`` holds, so it can be copied into a container whole. +the defaults) or ``ca.pem`` (the CA alone, for trust stores that add to the defaults). Each CA's +two files sit alone in a directory of their own under ``tls/trust/``, which can be copied into a +container whole. """ from __future__ import annotations +import functools import ipaddress import os import secrets @@ -38,7 +40,6 @@ _CA_LIFETIME = timedelta(days=5 * 365) _CA_RENEW_BEFORE = timedelta(days=30) # replaced this long before it expires -_SERVER_LIFETIME = timedelta(days=90) _CLOCK_SKEW = timedelta(minutes=5) _CA_FILE = "ca.key.pem" # the CA's key and certificate, private to the user @@ -76,8 +77,14 @@ def check_local_host(host: str) -> None: def local_ca() -> LocalCA: - """This user's local CA, created on first use and replaced when it nears expiry.""" - tls_dir = state_root() / "tls" + """This user's local CA, created on first use and replaced when it nears expiry. A process keeps the CA + it first loads, so its grant server and the containers it starts agree on it whatever another process + does to the one on disk.""" + return _process_ca(state_root() / "tls") + + +@functools.cache +def _process_ca(tls_dir: Path) -> LocalCA: ensure_state_dir(tls_dir) ca_file = tls_dir / _CA_FILE loaded = _load(ca_file) @@ -99,8 +106,9 @@ def local_ca() -> LocalCA: staged.unlink(missing_ok=True) loaded = _load(ca_file) or (key, cert) key, cert = loaded - trust_dir = tls_dir / _TRUST_DIR - trust_dir.mkdir(exist_ok=True) + # Written once per CA, so renewing the CA never changes the files another process's containers trust. + trust_dir = tls_dir / _TRUST_DIR / cert.fingerprint(hashes.SHA256()).hex()[:16] + trust_dir.mkdir(parents=True, exist_ok=True) trust_dir.chmod(0o755) cert_pem = cert.public_bytes(serialization.Encoding.PEM) cert_path = _write_if_changed(trust_dir / _CERT_FILE, cert_pem) @@ -120,7 +128,7 @@ def server_context(ca: LocalCA, hosts: Iterable[str]) -> ssl.SSLContext: .public_key(key.public_key()) .serial_number(x509.random_serial_number()) .not_valid_before(now - _CLOCK_SKEW) - .not_valid_after(min(now + _SERVER_LIFETIME, ca.cert.not_valid_after_utc)) + .not_valid_after(ca.cert.not_valid_after_utc) # a server keeps one certificate for its whole life .add_extension(x509.SubjectAlternativeName([_general_name(n) for n in names]), critical=False) .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True) .add_extension(_key_usage(digital_signature=True), critical=True) diff --git a/tst/unit/store/local_grants_server_test.py b/tst/unit/store/local_grants_server_test.py index fd53b1db..7a517ba0 100644 --- a/tst/unit/store/local_grants_server_test.py +++ b/tst/unit/store/local_grants_server_test.py @@ -169,6 +169,17 @@ def test_an_upload_larger_than_the_grant_is_refused(store): assert not any((store.root / ".agentenv-tmp").iterdir()) +def test_an_upload_with_oversized_part_headers_is_refused(store): + write = _policy(store) + body = ( + b"--b\r\nContent-Disposition: form-data; name=\"key\"\r\nX-Padding: " + b"x" * 20_000 + + b"\r\n\r\nns/x\r\n--b--\r\n" + ) + response = httpx.post(write.url, content=body, headers={"Content-Type": "multipart/form-data; boundary=b"}) + assert response.status_code == 400 + assert not store.exists("ns/x") + + def test_an_upload_must_name_its_key_before_its_file(store): write = _policy(store) assert _post(write, {"key": "ns/late"}, file_first=True).status_code == 400 diff --git a/tst/unit/store/local_grants_tls_test.py b/tst/unit/store/local_grants_tls_test.py index d209f443..a1282984 100644 --- a/tst/unit/store/local_grants_tls_test.py +++ b/tst/unit/store/local_grants_tls_test.py @@ -77,22 +77,57 @@ def test_the_ca_key_is_private_to_the_user_and_the_trust_files_are_not(): def test_the_ca_is_created_once_and_reused(): - assert local_ca().cert == local_ca().cert + first = local_ca() + tls._process_ca.cache_clear() # as another process would + assert local_ca().cert == first.cert + + +def test_a_process_keeps_its_ca_and_trust_files_when_another_renews_it(): + mine = local_ca() + key, cert = tls._new_ca() + (state_root() / "tls" / "ca.key.pem").write_bytes(tls._key_pem(key) + cert.public_bytes(tls.serialization.Encoding.PEM)) + + assert local_ca() is mine + assert x509.load_pem_x509_certificate(mine.cert_path.read_bytes()) == mine.cert + tls._process_ca.cache_clear() + theirs = local_ca() + assert theirs.cert == cert and theirs.trust_dir != mine.trust_dir + assert x509.load_pem_x509_certificate(mine.cert_path.read_bytes()) == mine.cert + + +def test_a_server_certificate_lasts_as_long_as_its_ca(): + ca = local_ca() + server = server_context(ca, ["localhost"]) + client = ssl.create_default_context(cafile=str(ca.cert_path)) + to_server, to_client = ssl.MemoryBIO(), ssl.MemoryBIO() + tls_client = client.wrap_bio(to_client, to_server, server_hostname="localhost") + tls_server = server.wrap_bio(to_server, to_client, server_side=True) + for _ in range(10): + for side in (tls_client, tls_server): + try: + side.do_handshake() + except ssl.SSLWantReadError: + pass + leaf = x509.load_der_x509_certificate(tls_client.getpeercert(binary_form=True)) + assert leaf.not_valid_after_utc == ca.cert.not_valid_after_utc def test_a_ca_near_expiry_is_replaced(monkeypatch): with monkeypatch.context() as m: m.setattr(tls, "_CA_LIFETIME", timedelta(days=10)) first = local_ca() + tls._process_ca.cache_clear() second = local_ca() assert second.cert != first.cert assert x509.load_pem_x509_certificate(second.cert_path.read_bytes()) == second.cert + tls._process_ca.cache_clear() assert local_ca().cert == second.cert def test_an_unreadable_ca_is_replaced(): first = local_ca() (state_root() / "tls" / "ca.key.pem").write_text("not a key") + tls._process_ca.cache_clear() assert local_ca().cert != first.cert From 1271813d061ea456344b5926a90ebfff6f010109 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 07:58:30 -1000 Subject: [PATCH 06/14] fix(store): a write-once put stays exclusive on a filesystem without hard links Where no hard link can claim the key, a write-once put now claims it with a marker created exclusively in the staging directory, held only from the existence check to the rename. A racing writer waits for the marker to go and then finds the winner's object; a marker a crash left behind is taken over once it is ten seconds old. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 35 ++++++++++++++-- tst/unit/store/local_object_store_test.py | 42 +++++++++++++++++++ 2 files changed, 73 insertions(+), 4 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index d6e10486..a1608648 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -4,11 +4,13 @@ import contextlib import errno +import hashlib import json import logging import os import shutil import tempfile +import time from collections.abc import Iterator from datetime import datetime, timezone from pathlib import Path @@ -23,8 +25,9 @@ _META_DIR = ".agentenv-meta" # each object's content type, at the object's own key _STAGING_DIR = ".agentenv-tmp" # files being written, inside the root so a rename into place is atomic _RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) -# A filesystem without hard links: a no-overwrite write checks, then renames, so two racing writers can both land. +# A filesystem without hard links: a no-overwrite write claims its key with a marker file instead. _NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) +_STALE_CLAIM_SECONDS = 10 # a claim is held only from the check to the rename; one this old was left by a crash class LocalFilesystemObjectStore(ObjectStore): @@ -182,9 +185,10 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit except OSError as e: if e.errno not in _NO_HARD_LINKS: raise - if path.exists(): - raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None - os.replace(staged, path) + with self._claimed(path): + if path.exists(): + raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None + os.replace(staged, path) # The default type is what a writer that named none gets; it reads back as unknown, so readers # still guess from the name as they did before types were kept. try: @@ -199,6 +203,29 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit except OSError as e: # the object is in place and its type reads back as unknown logger.warning("Could not record the content type of %s: %s", path, e) + @contextlib.contextmanager + def _claimed(self, path: Path) -> Iterator[None]: + """The sole right to create ``path``, where no hard link can claim it: a marker in the staging directory, + created exclusively. A writer that finds one waits for it to go, taking it over once it is stale.""" + marker = self._root / _STAGING_DIR / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" + while True: + try: + os.close(os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY)) + break + except FileExistsError: + try: + stale = time.time() - marker.stat().st_mtime > _STALE_CLAIM_SECONDS + except FileNotFoundError: + continue + if stale: + marker.unlink(missing_ok=True) + else: + time.sleep(0.01) + try: + yield + finally: + marker.unlink(missing_ok=True) + def _read_content_type(self, path: Path, st: os.stat_result | None = None) -> str | None: """The content type recorded for the object at ``path``, if it was recorded for this very file; None otherwise, as for one written before types were kept.""" diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index e8295c63..9002fdf9 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -5,7 +5,10 @@ """ import errno +import hashlib import os +import threading +import time from pathlib import Path import pytest @@ -190,6 +193,45 @@ def failing_replace(src, dst): assert not store.exists("nolink/y") +def test_without_hard_links_a_racing_write_once_put_keeps_the_winner(store, monkeypatch): + def no_links(src, dst): + raise OSError(errno.EPERM, "hard links not supported") + + real_replace = os.replace + first_holds_its_claim = threading.Event() + + def slow_first_replace(src, dst): + if threading.current_thread().name == "first" and str(dst).endswith("race/once"): + first_holds_its_claim.set() + time.sleep(0.3) + real_replace(src, dst) + + monkeypatch.setattr(os, "link", no_links) + monkeypatch.setattr(os, "replace", slow_first_replace) + first = threading.Thread(target=store.put, args=("race/once", b"first"), name="first") + first.start() + first_holds_its_claim.wait(5) + with pytest.raises(ObjectAlreadyExistsError): + store.put("race/once", b"second") + first.join() + assert store.read("race/once") == b"first" + + +def test_without_hard_links_a_claim_a_crash_left_is_taken_over(store, tmp_path, monkeypatch): + def no_links(src, dst): + raise OSError(errno.EPERM, "hard links not supported") + + monkeypatch.setattr(os, "link", no_links) + store.put("seed", b"s") # creates the staging directory + path = (tmp_path / "left/x").resolve() + marker = tmp_path / ".agentenv-tmp" / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" + marker.write_bytes(b"") + os.utime(marker, (time.time() - 60, time.time() - 60)) + store.put("left/x", b"mine") + assert store.read("left/x") == b"mine" + assert not marker.exists() + + def test_a_type_recorded_for_another_write_reads_back_as_unknown(store, tmp_path): """Two overwrites of one key race: the one that renamed first records its type last. Its type must not be read as the other's.""" From 589642321a42cff253dfceb1ca6c31d8a712255e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 08:14:20 -1000 Subject: [PATCH 07/14] fix(store): a live writer keeps its claim, and recorded types survive a copy that keeps modification times Without hard links, a write-once put's claim marker now names its owner's pid and is taken over only once that process has exited, not after ten seconds, so a slow live writer never loses it. On Windows, where rename refuses an existing target, the rename itself is the claim. A recorded content type is now paired with the file's size and modification time, and each write stamps its staged file with a modification time no other write shares (now, to the millisecond, plus a random remainder). Racing writes still never pair bytes with another write's type, and a copy of the store that keeps modification times keeps its types, where the inode it no longer paired with would not survive one. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 63 ++++++++++++++----- tst/unit/store/local_object_store_test.py | 41 ++++++++++-- 2 files changed, 84 insertions(+), 20 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index a1608648..77f1f6c7 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -8,6 +8,7 @@ import json import logging import os +import secrets import shutil import tempfile import time @@ -27,7 +28,7 @@ _RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) # A filesystem without hard links: a no-overwrite write claims its key with a marker file instead. _NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) -_STALE_CLAIM_SECONDS = 10 # a claim is held only from the check to the rename; one this old was left by a crash +_UNWRITTEN_CLAIM_SECONDS = 10 # a marker is given its owner's pid as it is created; one still empty this long never will be class LocalFilesystemObjectStore(ObjectStore): @@ -173,7 +174,7 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit The type is recorded with the identity of the file it describes, so one recorded for another write (a racing overwrite, or one whose type could not be recorded) reads back as unknown.""" - identity = _identity(staged.stat()) + identity = _stamp(staged) path.parent.mkdir(parents=True, exist_ok=True) if allow_overwrite: os.replace(staged, path) @@ -185,10 +186,16 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit except OSError as e: if e.errno not in _NO_HARD_LINKS: raise - with self._claimed(path): - if path.exists(): + if os.name == "nt": # rename there refuses an existing target, which is the claim + try: + os.rename(staged, path) + except FileExistsError: raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None - os.replace(staged, path) + else: + with self._claimed(path): + if path.exists(): + raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None + os.replace(staged, path) # The default type is what a writer that named none gets; it reads back as unknown, so readers # still guess from the name as they did before types were kept. try: @@ -206,21 +213,21 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit @contextlib.contextmanager def _claimed(self, path: Path) -> Iterator[None]: """The sole right to create ``path``, where no hard link can claim it: a marker in the staging directory, - created exclusively. A writer that finds one waits for it to go, taking it over once it is stale.""" + created exclusively and naming its owner's pid. A writer that finds one waits for it to go, and takes it + over only once its owner has exited.""" marker = self._root / _STAGING_DIR / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" while True: try: - os.close(os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY)) - break + fd = os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY) except FileExistsError: - try: - stale = time.time() - marker.stat().st_mtime > _STALE_CLAIM_SECONDS - except FileNotFoundError: - continue - if stale: + if _abandoned(marker): marker.unlink(missing_ok=True) else: time.sleep(0.01) + continue + with os.fdopen(fd, "w") as f: + f.write(str(os.getpid())) + break try: yield finally: @@ -250,9 +257,35 @@ def _from_url(object_url: str) -> str: return object_url[len("file://"):] if object_url.startswith("file://") else object_url +def _stamp(staged: Path) -> list[int]: + """Give a staged file a modification time no other write shares (now, to the millisecond, plus a random + remainder) and return its identity. A rename keeps it, and so does a copy that keeps modification times.""" + now_ns = time.time_ns() + stamp = now_ns - now_ns % 1_000_000 + secrets.randbelow(1_000_000) + os.utime(staged, ns=(stamp, stamp)) + return _identity(staged.stat()) + + def _identity(st: os.stat_result) -> list[int]: - """What tells one file apart from another written to the same key: a rename keeps all three.""" - return [st.st_ino, st.st_size, st.st_mtime_ns] + return [st.st_size, st.st_mtime_ns] + + +def _abandoned(marker: Path) -> bool: + """Whether a claim marker's owner has exited, or never got to name itself.""" + try: + owner = marker.read_text() + age = time.time() - marker.stat().st_mtime + except FileNotFoundError: + return False + if not owner: + return age > _UNWRITTEN_CLAIM_SECONDS + try: + os.kill(int(owner), 0) + except ProcessLookupError: + return True + except (PermissionError, ValueError): + return False + return False @contextlib.contextmanager diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index 9002fdf9..860e6f78 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -7,6 +7,9 @@ import errno import hashlib import os +import shutil +import subprocess +import sys import threading import time from pathlib import Path @@ -217,21 +220,49 @@ def slow_first_replace(src, dst): assert store.read("race/once") == b"first" -def test_without_hard_links_a_claim_a_crash_left_is_taken_over(store, tmp_path, monkeypatch): +def _claim_marker(tmp_path, key): + path = (tmp_path / key).resolve() + return tmp_path / ".agentenv-tmp" / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" + + +@pytest.fixture +def no_hard_links(store, monkeypatch): def no_links(src, dst): raise OSError(errno.EPERM, "hard links not supported") monkeypatch.setattr(os, "link", no_links) store.put("seed", b"s") # creates the staging directory - path = (tmp_path / "left/x").resolve() - marker = tmp_path / ".agentenv-tmp" / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" - marker.write_bytes(b"") - os.utime(marker, (time.time() - 60, time.time() - 60)) + + +def test_without_hard_links_a_claim_whose_owner_exited_is_taken_over(store, tmp_path, no_hard_links): + exited = subprocess.run([sys.executable, "-c", "import os; print(os.getpid())"], capture_output=True, text=True) + marker = _claim_marker(tmp_path, "left/x") + marker.write_text(exited.stdout.strip()) store.put("left/x", b"mine") assert store.read("left/x") == b"mine" assert not marker.exists() +def test_without_hard_links_a_live_writers_claim_is_never_taken_over(store, tmp_path, no_hard_links): + marker = _claim_marker(tmp_path, "held/x") + marker.write_text(str(os.getpid())) + os.utime(marker, (time.time() - 3600, time.time() - 3600)) + writer = threading.Thread(target=store.put, args=("held/x", b"late")) + writer.start() + writer.join(0.3) + assert writer.is_alive() and not store.exists("held/x") + marker.unlink() + writer.join(5) + assert store.read("held/x") == b"late" + + +def test_types_survive_a_copy_that_keeps_modification_times(store, tmp_path): + store.put("copied/x", b"typed", content_type="text/plain") + copy = tmp_path.parent / f"{tmp_path.name}-copy" + shutil.copytree(tmp_path, copy) # copies modification times, not inodes + assert LocalFilesystemObjectStore(str(copy)).get_object_metadata("copied/x").content_type == "text/plain" + + def test_a_type_recorded_for_another_write_reads_back_as_unknown(store, tmp_path): """Two overwrites of one key race: the one that renamed first records its type last. Its type must not be read as the other's.""" From a2d2793bac46039015dc4945953fcf7d04f83312 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 08:24:04 -1000 Subject: [PATCH 08/14] fix(store): claim a key without hard links with a lock that dies with its holder A write-once put on a filesystem without hard links now claims its key with an exclusive flock on a file in the staging directory rather than a marker naming a pid. The kernel drops the lock when its holder exits, so a crashed writer never blocks the key and a reused pid can never pass for a live one; a live holder is waited for however long it takes. The file stays in place so every writer locks the same one. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 49 +++++-------------- tst/unit/store/local_object_store_test.py | 38 +++++++------- 2 files changed, 32 insertions(+), 55 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index 77f1f6c7..a60cf3c7 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -18,6 +18,9 @@ from typing import BinaryIO from agent_env.store.base import ObjectAlreadyExistsError, ObjectNotFoundError + +if os.name != "nt": + import fcntl from agent_env.store.local_state import ensure_state_dir from agent_env.store.object_store.object_store import DEFAULT_CONTENT_TYPE, ObjectMetadata, ObjectStore @@ -26,9 +29,8 @@ _META_DIR = ".agentenv-meta" # each object's content type, at the object's own key _STAGING_DIR = ".agentenv-tmp" # files being written, inside the root so a rename into place is atomic _RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) -# A filesystem without hard links: a no-overwrite write claims its key with a marker file instead. +# A filesystem without hard links: a no-overwrite write claims its key with a lock instead. _NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) -_UNWRITTEN_CLAIM_SECONDS = 10 # a marker is given its owner's pid as it is created; one still empty this long never will be class LocalFilesystemObjectStore(ObjectStore): @@ -212,26 +214,16 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit @contextlib.contextmanager def _claimed(self, path: Path) -> Iterator[None]: - """The sole right to create ``path``, where no hard link can claim it: a marker in the staging directory, - created exclusively and naming its owner's pid. A writer that finds one waits for it to go, and takes it - over only once its owner has exited.""" + """The sole right to create ``path``, where no hard link can claim it: an exclusive lock on a file in the + staging directory, which the kernel drops if its holder dies. The file stays, so every writer locks the + same one.""" marker = self._root / _STAGING_DIR / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" - while True: + with marker.open("a") as f: + fcntl.flock(f, fcntl.LOCK_EX) try: - fd = os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY) - except FileExistsError: - if _abandoned(marker): - marker.unlink(missing_ok=True) - else: - time.sleep(0.01) - continue - with os.fdopen(fd, "w") as f: - f.write(str(os.getpid())) - break - try: - yield - finally: - marker.unlink(missing_ok=True) + yield + finally: + fcntl.flock(f, fcntl.LOCK_UN) def _read_content_type(self, path: Path, st: os.stat_result | None = None) -> str | None: """The content type recorded for the object at ``path``, if it was recorded for this very file; None @@ -270,23 +262,6 @@ def _identity(st: os.stat_result) -> list[int]: return [st.st_size, st.st_mtime_ns] -def _abandoned(marker: Path) -> bool: - """Whether a claim marker's owner has exited, or never got to name itself.""" - try: - owner = marker.read_text() - age = time.time() - marker.stat().st_mtime - except FileNotFoundError: - return False - if not owner: - return age > _UNWRITTEN_CLAIM_SECONDS - try: - os.kill(int(owner), 0) - except ProcessLookupError: - return True - except (PermissionError, ValueError): - return False - return False - @contextlib.contextmanager def _missing_as_not_found(object_url: str) -> Iterator[None]: diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index 860e6f78..cdbe38e5 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -234,26 +234,28 @@ def no_links(src, dst): store.put("seed", b"s") # creates the staging directory -def test_without_hard_links_a_claim_whose_owner_exited_is_taken_over(store, tmp_path, no_hard_links): - exited = subprocess.run([sys.executable, "-c", "import os; print(os.getpid())"], capture_output=True, text=True) +def test_without_hard_links_a_claim_dies_with_its_holder(store, tmp_path, no_hard_links): + """The claim is a lock the kernel drops when its holder exits, so a crashed writer never blocks the key.""" marker = _claim_marker(tmp_path, "left/x") - marker.write_text(exited.stdout.strip()) - store.put("left/x", b"mine") + holder = subprocess.Popen( + [sys.executable, "-c", ( + "import fcntl, sys, time; f = open(sys.argv[1], 'a'); fcntl.flock(f, fcntl.LOCK_EX); " + "print('locked', flush=True); time.sleep(60)" + ), str(marker)], + stdout=subprocess.PIPE, text=True, + ) + try: + assert holder.stdout.readline().strip() == "locked" + writer = threading.Thread(target=store.put, args=("left/x", b"mine")) + writer.start() + writer.join(0.3) + assert writer.is_alive() and not store.exists("left/x") # a live holder, however old, is waited for + holder.kill() + holder.wait() + writer.join(5) + finally: + holder.kill() assert store.read("left/x") == b"mine" - assert not marker.exists() - - -def test_without_hard_links_a_live_writers_claim_is_never_taken_over(store, tmp_path, no_hard_links): - marker = _claim_marker(tmp_path, "held/x") - marker.write_text(str(os.getpid())) - os.utime(marker, (time.time() - 3600, time.time() - 3600)) - writer = threading.Thread(target=store.put, args=("held/x", b"late")) - writer.start() - writer.join(0.3) - assert writer.is_alive() and not store.exists("held/x") - marker.unlink() - writer.join(5) - assert store.read("held/x") == b"late" def test_types_survive_a_copy_that_keeps_modification_times(store, tmp_path): From d46296db1578b40c40be6eb43bd2ef683914123d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 09:50:18 -1000 Subject: [PATCH 09/14] fix(store): a key's bytes and type are written and read under its lock, and staged files of dead writes are swept Chaos testing on exFAT, which keeps modification times to 10 ms, showed two same-size overwrites of one key could leave one's type recorded against the other's bytes (5 in 4,129 races): the per-write time stamp that tells writes apart is rounded away there. A write now holds the key's lock (an exclusive flock) across placing its bytes and recording its type, and a metadata read holds it shared, so no write lands between another's two steps and no reader sees one write's type with another's bytes, whatever the filesystem's time resolution. An overwrite drops the old type first, so one cut short leaves the type unknown, never wrong. Keys share 64 lock files, which also replace the per-key claim files the no-hard-link path used; Windows, which has no flock, keeps relying on the stamp, which NTFS's 100 ns times preserve. A store's first write also removes files staged by writes that died (untouched for an hour), which crashes left behind in .agentenv-tmp. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 132 +++++++++++------- tst/unit/store/local_object_store_test.py | 76 ++++++++-- 2 files changed, 143 insertions(+), 65 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index a60cf3c7..d15aab29 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -29,8 +29,11 @@ _META_DIR = ".agentenv-meta" # each object's content type, at the object's own key _STAGING_DIR = ".agentenv-tmp" # files being written, inside the root so a rename into place is atomic _RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) -# A filesystem without hard links: a no-overwrite write claims its key with a lock instead. +# A filesystem without hard links: a no-overwrite write relies on the key's lock instead. _NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) +_LOCK_STRIPES = 64 # keys share this many lock files, so locking leaves no file behind per key +_STAGED_PREFIX = "staged-" +_STALE_STAGED_SECONDS = 3600 # a staged file untouched this long belongs to a write that died class LocalFilesystemObjectStore(ObjectStore): @@ -41,6 +44,7 @@ class LocalFilesystemObjectStore(ObjectStore): def __init__(self, root: str) -> None: self._root = Path(root) + self._swept = False @property def root(self) -> Path: @@ -94,11 +98,13 @@ def open(self, object_url: str) -> BinaryIO: def get_object_metadata(self, key: str) -> ObjectMetadata | None: path = self._resolve(key) - if not path.is_file(): - return None - st = path.stat() + with self._locked(path, shared=True): # never between another write's bytes and its type + if not path.is_file(): + return None + st = path.stat() + content_type = self._read_content_type(path, st) return ObjectMetadata( - content_type=self._read_content_type(path, st), + content_type=content_type, size=st.st_size, last_modified=datetime.fromtimestamp(st.st_mtime, tz=timezone.utc), ) @@ -161,7 +167,10 @@ def _staged(self) -> Iterator[Path]: ensure_state_dir(self._root) staging = self._root / _STAGING_DIR staging.mkdir(exist_ok=True) - fd, name = tempfile.mkstemp(dir=staging) + if not self._swept: + self._swept = True + _sweep(staging) + fd, name = tempfile.mkstemp(dir=staging, prefix=_STAGED_PREFIX) os.close(fd) staged = Path(name) staged.chmod(0o644) @@ -171,59 +180,70 @@ def _staged(self) -> Iterator[Path]: staged.unlink(missing_ok=True) def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrite: bool) -> None: - """Put a staged file in place at ``path`` and record its content type. An overwrite renames over - whatever is there; otherwise a hard link claims ``path`` only if no other writer got there first. + """Put a staged file in place at ``path`` and record its content type, both under the key's lock, so no + other write to the key lands in between and no reader sees one write's type with another's bytes. An + overwrite drops the old type first, so one cut short leaves the type unknown, never wrong. - The type is recorded with the identity of the file it describes, so one recorded for another - write (a racing overwrite, or one whose type could not be recorded) reads back as unknown.""" + The type is also recorded with the identity of the file it describes, which is what keeps them paired + where there is no lock (Windows, whose NTFS keeps times fine enough to tell writes apart).""" identity = _stamp(staged) + meta = self._meta_path(path) path.parent.mkdir(parents=True, exist_ok=True) - if allow_overwrite: - os.replace(staged, path) - else: - try: - os.link(staged, path) - except FileExistsError: - raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None - except OSError as e: - if e.errno not in _NO_HARD_LINKS: - raise - if os.name == "nt": # rename there refuses an existing target, which is the claim + with self._locked(path): + if allow_overwrite: + meta.unlink(missing_ok=True) + os.replace(staged, path) + else: + try: + os.link(staged, path) # exclusive even against a writer that takes no lock + except FileExistsError: + raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None + except OSError as e: + if e.errno not in _NO_HARD_LINKS: + raise + # The lock makes the check and the rename one step; on Windows, rename refuses an existing target. + if path.exists(): + raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None try: - os.rename(staged, path) + (os.rename if os.name == "nt" else os.replace)(staged, path) except FileExistsError: raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None - else: - with self._claimed(path): - if path.exists(): - raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None - os.replace(staged, path) - # The default type is what a writer that named none gets; it reads back as unknown, so readers - # still guess from the name as they did before types were kept. - try: - meta = self._meta_path(path) - if content_type == DEFAULT_CONTENT_TYPE: - meta.unlink(missing_ok=True) - return - meta.parent.mkdir(parents=True, exist_ok=True) - with self._staged() as staged_meta: - staged_meta.write_text(json.dumps({"content_type": content_type, "object": identity})) - os.replace(staged_meta, meta) - except OSError as e: # the object is in place and its type reads back as unknown - logger.warning("Could not record the content type of %s: %s", path, e) + # The default type is what a writer that named none gets; it reads back as unknown, so readers + # still guess from the name as they did before types were kept. + try: + if content_type == DEFAULT_CONTENT_TYPE: + meta.unlink(missing_ok=True) + return + meta.parent.mkdir(parents=True, exist_ok=True) + with self._staged() as staged_meta: + staged_meta.write_text(json.dumps({"content_type": content_type, "object": identity})) + os.replace(staged_meta, meta) + except OSError as e: # the object is in place and its type reads back as unknown + logger.warning("Could not record the content type of %s: %s", path, e) @contextlib.contextmanager - def _claimed(self, path: Path) -> Iterator[None]: - """The sole right to create ``path``, where no hard link can claim it: an exclusive lock on a file in the - staging directory, which the kernel drops if its holder dies. The file stays, so every writer locks the - same one.""" - marker = self._root / _STAGING_DIR / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" - with marker.open("a") as f: - fcntl.flock(f, fcntl.LOCK_EX) - try: - yield - finally: - fcntl.flock(f, fcntl.LOCK_UN) + def _locked(self, path: Path, *, shared: bool = False) -> Iterator[None]: + """The key's lock, exclusive to write it and shared to read its type: a flock on one of a fixed set of + files in the staging directory, which the kernel drops if its holder dies. None on Windows.""" + if os.name == "nt": + yield + return + try: + f = self._lock_path(path).open("r" if shared else "a") + except OSError: + if not shared: + raise + f = None # no write has taken this lock yet, or the store is read-only: nothing to wait for + if f is None: + yield + return + with f: + fcntl.flock(f, fcntl.LOCK_SH if shared else fcntl.LOCK_EX) + yield + + def _lock_path(self, path: Path) -> Path: + stripe = int(hashlib.sha256(str(path).encode()).hexdigest()[:8], 16) % _LOCK_STRIPES + return self._root / _STAGING_DIR / f"lock-{stripe:02x}" def _read_content_type(self, path: Path, st: os.stat_result | None = None) -> str | None: """The content type recorded for the object at ``path``, if it was recorded for this very file; None @@ -258,6 +278,16 @@ def _stamp(staged: Path) -> list[int]: return _identity(staged.stat()) +def _sweep(staging: Path) -> None: + """Remove files staged by writes that died. A write in progress keeps touching its file, and a finished one + renames it away, so one untouched for an hour will never be committed.""" + cutoff = time.time() - _STALE_STAGED_SECONDS + for staged in staging.glob(f"{_STAGED_PREFIX}*"): + with contextlib.suppress(OSError): + if staged.stat().st_mtime < cutoff: + staged.unlink() + + def _identity(st: os.stat_result) -> list[int]: return [st.st_size, st.st_mtime_ns] diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index cdbe38e5..72213a13 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -220,9 +220,17 @@ def slow_first_replace(src, dst): assert store.read("race/once") == b"first" -def _claim_marker(tmp_path, key): - path = (tmp_path / key).resolve() - return tmp_path / ".agentenv-tmp" / f"{hashlib.sha256(str(path).encode()).hexdigest()}.claim" +def _hold_lock(lock_file): + """A process holding ``lock_file`` exclusively until killed, as a writer mid-commit would.""" + holder = subprocess.Popen( + [sys.executable, "-c", ( + "import fcntl, sys, time; f = open(sys.argv[1], 'a'); fcntl.flock(f, fcntl.LOCK_EX); " + "print('locked', flush=True); time.sleep(60)" + ), str(lock_file)], + stdout=subprocess.PIPE, text=True, + ) + assert holder.stdout.readline().strip() == "locked" + return holder @pytest.fixture @@ -236,20 +244,12 @@ def no_links(src, dst): def test_without_hard_links_a_claim_dies_with_its_holder(store, tmp_path, no_hard_links): """The claim is a lock the kernel drops when its holder exits, so a crashed writer never blocks the key.""" - marker = _claim_marker(tmp_path, "left/x") - holder = subprocess.Popen( - [sys.executable, "-c", ( - "import fcntl, sys, time; f = open(sys.argv[1], 'a'); fcntl.flock(f, fcntl.LOCK_EX); " - "print('locked', flush=True); time.sleep(60)" - ), str(marker)], - stdout=subprocess.PIPE, text=True, - ) + holder = _hold_lock(store._lock_path((tmp_path / "left/x").resolve())) try: - assert holder.stdout.readline().strip() == "locked" writer = threading.Thread(target=store.put, args=("left/x", b"mine")) writer.start() writer.join(0.3) - assert writer.is_alive() and not store.exists("left/x") # a live holder, however old, is waited for + assert writer.is_alive() and not (tmp_path / "left/x").exists() # a live holder is waited for holder.kill() holder.wait() writer.join(5) @@ -295,7 +295,55 @@ def test_listings_leave_out_the_stores_own_files(store, tmp_path): (tmp_path / ".agentenv-tmp" / "crash-residue").write_bytes(b"half") assert store.list("") == ["a/x.json"] assert store.list_at(store.object_url("")) == [store.object_url("a/x.json")] - assert os.listdir(tmp_path / ".agentenv-tmp") == ["crash-residue"] # committed writes leave nothing staged + staging = [name for name in os.listdir(tmp_path / ".agentenv-tmp") if not name.startswith("lock-")] + assert staging == ["crash-residue"] # committed writes leave nothing staged + + +def test_a_type_is_never_read_while_a_write_to_the_key_is_in_progress(store, tmp_path): + store.put("busy/x", b"typed", content_type="text/plain") + holder = _hold_lock(store._lock_path((tmp_path / "busy/x").resolve())) + try: + read = [] + reader = threading.Thread(target=lambda: read.append(store.get_object_metadata("busy/x"))) + reader.start() + reader.join(0.3) + assert reader.is_alive() # waits for the write holding the key + holder.kill() + holder.wait() + reader.join(5) + finally: + holder.kill() + assert read[0].content_type == "text/plain" + + +def test_an_overwrite_cut_short_leaves_the_type_unknown_not_the_old_one(store, tmp_path, monkeypatch): + store.put("cut/x", b"first", content_type="text/plain") + real_replace = os.replace + + def replace(src, dst): + if str(dst).endswith("cut/x") and ".agentenv-meta" not in str(dst): + raise OSError(errno.EIO, "disk went away") + real_replace(src, dst) + + monkeypatch.setattr(os, "replace", replace) + with pytest.raises(OSError, match="disk went away"): + store.put("cut/x", b"second", content_type="application/json", allow_overwrite=True) + assert (store.read("cut/x"), store.get_object_metadata("cut/x").content_type) == (b"first", None) + + +def test_a_store_sweeps_files_staged_by_writes_that_died(tmp_path): + staging = tmp_path / ".agentenv-tmp" + staging.mkdir() + dead, live = staging / "staged-dead", staging / "staged-live" + dead.write_bytes(b"half") + live.write_bytes(b"in progress") + os.utime(dead, (time.time() - 7200, time.time() - 7200)) + (staging / "lock-00").write_bytes(b"") + os.utime(staging / "lock-00", (time.time() - 7200, time.time() - 7200)) + + LocalFilesystemObjectStore(str(tmp_path)).put("k", b"v") + + assert not dead.exists() and live.exists() and (staging / "lock-00").exists() def test_list_of_an_unwritten_store_is_empty(tmp_path): From 7bf3aa2d5be00b293836529185fec5bfc080975a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 10:23:07 -1000 Subject: [PATCH 10/14] fix(store): a grant download sends the bytes and type of one write The GET handler now opens the object under the key's shared lock and reads its type against that open file, then streams from it, so a racing overwrite can no longer pair its bytes with the type of the object that was there when the type was read. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_grants/server.py | 28 ++++++++++++++---- tst/unit/store/local_grants_server_test.py | 29 +++++++++++++++++++ 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/agent_env/store/object_store/local_grants/server.py b/src/agent_env/store/object_store/local_grants/server.py index 56ba81f0..57605fb1 100644 --- a/src/agent_env/store/object_store/local_grants/server.py +++ b/src/agent_env/store/object_store/local_grants/server.py @@ -13,12 +13,14 @@ import functools import ipaddress import logging +import os import platform import socket import subprocess import threading import time from pathlib import Path +from collections.abc import AsyncIterator from typing import TYPE_CHECKING, BinaryIO import uvicorn @@ -26,7 +28,7 @@ from python_multipart.multipart import MultipartParser, parse_options_header from starlette.applications import Starlette from starlette.requests import ClientDisconnect, Request -from starlette.responses import FileResponse, PlainTextResponse, Response +from starlette.responses import PlainTextResponse, Response, StreamingResponse from starlette.routing import Route from uvicorn.protocols.http.h11_impl import H11Protocol @@ -54,6 +56,7 @@ _MAX_FORM_FIELDS = 32 _MAX_FORM_FIELD_BYTES = 64 * 1024 _MAX_PART_HEADER_BYTES = 16 * 1024 +_DOWNLOAD_CHUNK_BYTES = 1024 * 1024 _servers: dict[tuple[str, str], GrantServer] = {} _servers_lock = threading.Lock() @@ -238,10 +241,25 @@ def _path(store: LocalFilesystemObjectStore, key: str) -> Path: def _get(store: LocalFilesystemObjectStore, claims: GrantClaims) -> Response: path = _path(store, claims.key) - metadata = store.get_object_metadata(claims.key) # read under the key's lock, never mid-write - if metadata is None: - raise _Rejected(404, "No object exists at this grant's key.") - return FileResponse(path, media_type=metadata.content_type or DEFAULT_CONTENT_TYPE) + # Opened under the key's lock and typed from that open file, so the bytes sent and their type are one write's. + with store._locked(path, shared=True): + try: + f = path.open("rb") + except (FileNotFoundError, IsADirectoryError, NotADirectoryError): + raise _Rejected(404, "No object exists at this grant's key.") from None + st = os.fstat(f.fileno()) + content_type = store._read_content_type(path, st) + return StreamingResponse( + _chunks(f), media_type=content_type or DEFAULT_CONTENT_TYPE, headers={"Content-Length": str(st.st_size)} + ) + + +async def _chunks(f: BinaryIO) -> AsyncIterator[bytes]: + try: + while chunk := await asyncio.to_thread(f.read, _DOWNLOAD_CHUNK_BYTES): + yield chunk + finally: + f.close() async def _put(store: LocalFilesystemObjectStore, claims: GrantClaims, request: Request) -> None: diff --git a/tst/unit/store/local_grants_server_test.py b/tst/unit/store/local_grants_server_test.py index 7a517ba0..96ebdb38 100644 --- a/tst/unit/store/local_grants_server_test.py +++ b/tst/unit/store/local_grants_server_test.py @@ -1,7 +1,10 @@ """The local store's grants against its real grant server, over HTTPS on loopback: the conformance grant cases, and the bounds and refusals S3 would apply to the same requests.""" +import os import subprocess +import sys +import threading import time import httpx @@ -66,6 +69,32 @@ def test_a_read_returns_the_exact_bytes_and_type(store): assert "content-encoding" not in response.headers +@pytest.mark.skipif(os.name == "nt", reason="holds the key's lock with flock, which Windows lacks") +def test_a_download_waits_for_a_write_to_its_key(store): + """Bytes and type are read from one file under the key's lock, so a download never mixes two writes.""" + grant = store.issue_read_grant(store.put("a/held", b"typed", content_type="text/plain")).url + holder = subprocess.Popen( + [sys.executable, "-c", ( + "import fcntl, sys, time; f = open(sys.argv[1], 'a'); fcntl.flock(f, fcntl.LOCK_EX); " + "print('locked', flush=True); time.sleep(60)" + ), str(store._lock_path((store.root / "a/held").resolve()))], + stdout=subprocess.PIPE, text=True, + ) + try: + assert holder.stdout.readline().strip() == "locked" + got = [] + download = threading.Thread(target=lambda: got.append(httpx.get(grant, timeout=10))) + download.start() + download.join(0.3) + assert download.is_alive() + holder.kill() + holder.wait() + download.join(10) + finally: + holder.kill() + assert (got[0].status_code, got[0].content, got[0].headers["content-type"]) == (200, b"typed", "text/plain; charset=utf-8") + + def test_a_read_of_a_missing_object_is_404(store): assert httpx.get(store.issue_read_grant(store.object_url("a/none.bin")).url).status_code == 404 From 1b8b548d60936a8ba296155a61c86ddedc170f17 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 10:26:37 -1000 Subject: [PATCH 11/14] fix(store): a failed overwrite keeps the old type, and the sweep spares any file a live write holds An overwrite now sets the old type record aside while its bytes land and puts it back if they do not, so a rename that fails keeps the old object's type; one cut short after its bytes land still reads as unknown, never as the old type. Each staged file is held under an exclusive flock by its writer for as long as the write lives, and the sweep removes only staged files no process holds (and none younger than a minute, which may not be held yet), so a write paused for any length of time never loses its file. Windows, which has no flock, is not swept. Keys now share 256 lock files rather than 64, so fewer unrelated keys wait on one another, and the tests that hold a lock with flock skip on Windows. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 53 +++++++++++++++---- tst/unit/store/local_object_store_test.py | 50 +++++++++++------ 2 files changed, 77 insertions(+), 26 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index d15aab29..27dbf90c 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -31,9 +31,9 @@ _RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) # A filesystem without hard links: a no-overwrite write relies on the key's lock instead. _NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) -_LOCK_STRIPES = 64 # keys share this many lock files, so locking leaves no file behind per key +_LOCK_STRIPES = 256 # keys share this many lock files, so locking leaves no file behind per key _STAGED_PREFIX = "staged-" -_STALE_STAGED_SECONDS = 3600 # a staged file untouched this long belongs to a write that died +_STAGED_GRACE_SECONDS = 60 # a staged file younger than this may not be locked by its writer yet class LocalFilesystemObjectStore(ObjectStore): @@ -171,12 +171,17 @@ def _staged(self) -> Iterator[Path]: self._swept = True _sweep(staging) fd, name = tempfile.mkstemp(dir=staging, prefix=_STAGED_PREFIX) - os.close(fd) staged = Path(name) - staged.chmod(0o644) + if os.name == "nt": # an open file cannot be renamed there, and there is no flock to hold + os.close(fd) + else: + fcntl.flock(fd, fcntl.LOCK_EX) # held while the write lives, so a sweep never takes its file try: + staged.chmod(0o644) yield staged finally: + if os.name != "nt": + os.close(fd) staged.unlink(missing_ok=True) def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrite: bool) -> None: @@ -191,8 +196,18 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit path.parent.mkdir(parents=True, exist_ok=True) with self._locked(path): if allow_overwrite: - meta.unlink(missing_ok=True) - os.replace(staged, path) + # The old type is set aside while the bytes land: a write cut short reads as unknown, never as the + # old type, and one that fails before its bytes land puts the old type back. + aside = self._set_aside(meta) + try: + os.replace(staged, path) + except BaseException: + if aside is not None: + with contextlib.suppress(FileNotFoundError): # swept meanwhile: the type reads as unknown + os.replace(aside, meta) + raise + if aside is not None: + aside.unlink(missing_ok=True) else: try: os.link(staged, path) # exclusive even against a writer that takes no lock @@ -241,6 +256,17 @@ def _locked(self, path: Path, *, shared: bool = False) -> Iterator[None]: fcntl.flock(f, fcntl.LOCK_SH if shared else fcntl.LOCK_EX) yield + def _set_aside(self, meta: Path) -> Path | None: + """Move a type record into the staging directory, fresh enough that no sweep takes it; None when there is + none.""" + aside = self._root / _STAGING_DIR / f"{_STAGED_PREFIX}aside-{secrets.token_hex(8)}" + try: + os.replace(meta, aside) + except FileNotFoundError: + return None + os.utime(aside) + return aside + def _lock_path(self, path: Path) -> Path: stripe = int(hashlib.sha256(str(path).encode()).hexdigest()[:8], 16) % _LOCK_STRIPES return self._root / _STAGING_DIR / f"lock-{stripe:02x}" @@ -279,12 +305,17 @@ def _stamp(staged: Path) -> list[int]: def _sweep(staging: Path) -> None: - """Remove files staged by writes that died. A write in progress keeps touching its file, and a finished one - renames it away, so one untouched for an hour will never be committed.""" - cutoff = time.time() - _STALE_STAGED_SECONDS + """Remove files staged by writes that died: each live write holds a lock on its file, which the kernel drops + when the writer exits, so a file no one holds will never be committed. Not on Windows, which has no flock.""" + if os.name == "nt": + return + cutoff = time.time() - _STAGED_GRACE_SECONDS for staged in staging.glob(f"{_STAGED_PREFIX}*"): - with contextlib.suppress(OSError): - if staged.stat().st_mtime < cutoff: + with contextlib.suppress(OSError): # BlockingIOError, the one that matters: a live write holds it + if staged.stat().st_mtime > cutoff: + continue + with staged.open("rb") as f: + fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB) staged.unlink() diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index 72213a13..549c00fa 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -220,6 +220,9 @@ def slow_first_replace(src, dst): assert store.read("race/once") == b"first" +_needs_flock = pytest.mark.skipif(os.name == "nt", reason="holds a lock with flock, which Windows lacks") + + def _hold_lock(lock_file): """A process holding ``lock_file`` exclusively until killed, as a writer mid-commit would.""" holder = subprocess.Popen( @@ -242,6 +245,7 @@ def no_links(src, dst): store.put("seed", b"s") # creates the staging directory +@_needs_flock def test_without_hard_links_a_claim_dies_with_its_holder(store, tmp_path, no_hard_links): """The claim is a lock the kernel drops when its holder exits, so a crashed writer never blocks the key.""" holder = _hold_lock(store._lock_path((tmp_path / "left/x").resolve())) @@ -299,6 +303,7 @@ def test_listings_leave_out_the_stores_own_files(store, tmp_path): assert staging == ["crash-residue"] # committed writes leave nothing staged +@_needs_flock def test_a_type_is_never_read_while_a_write_to_the_key_is_in_progress(store, tmp_path): store.put("busy/x", b"typed", content_type="text/plain") holder = _hold_lock(store._lock_path((tmp_path / "busy/x").resolve())) @@ -316,34 +321,49 @@ def test_a_type_is_never_read_while_a_write_to_the_key_is_in_progress(store, tmp assert read[0].content_type == "text/plain" -def test_an_overwrite_cut_short_leaves_the_type_unknown_not_the_old_one(store, tmp_path, monkeypatch): - store.put("cut/x", b"first", content_type="text/plain") +def _failing_replace(monkeypatch, suffix): real_replace = os.replace def replace(src, dst): - if str(dst).endswith("cut/x") and ".agentenv-meta" not in str(dst): + if str(dst).endswith(suffix): raise OSError(errno.EIO, "disk went away") real_replace(src, dst) monkeypatch.setattr(os, "replace", replace) + + +def test_an_overwrite_whose_bytes_do_not_land_keeps_the_old_type(store, tmp_path, monkeypatch): + store.put("cut/x", b"first", content_type="text/plain") + _failing_replace(monkeypatch, str(tmp_path / "cut/x")) with pytest.raises(OSError, match="disk went away"): store.put("cut/x", b"second", content_type="application/json", allow_overwrite=True) - assert (store.read("cut/x"), store.get_object_metadata("cut/x").content_type) == (b"first", None) + assert (store.read("cut/x"), store.get_object_metadata("cut/x").content_type) == (b"first", "text/plain") -def test_a_store_sweeps_files_staged_by_writes_that_died(tmp_path): - staging = tmp_path / ".agentenv-tmp" - staging.mkdir() - dead, live = staging / "staged-dead", staging / "staged-live" - dead.write_bytes(b"half") - live.write_bytes(b"in progress") - os.utime(dead, (time.time() - 7200, time.time() - 7200)) - (staging / "lock-00").write_bytes(b"") - os.utime(staging / "lock-00", (time.time() - 7200, time.time() - 7200)) +def test_an_overwrite_whose_type_cannot_be_recorded_reads_unknown_not_the_old_type(store, tmp_path, monkeypatch): + store.put("cut/y", b"first", content_type="text/plain") + _failing_replace(monkeypatch, str(tmp_path / ".agentenv-meta" / "cut/y")) + store.put("cut/y", b"second!", content_type="application/json", allow_overwrite=True) + assert (store.read("cut/y"), store.get_object_metadata("cut/y").content_type) == (b"second!", None) - LocalFilesystemObjectStore(str(tmp_path)).put("k", b"v") - assert not dead.exists() and live.exists() and (staging / "lock-00").exists() +@_needs_flock +def test_a_store_sweeps_only_files_no_live_write_holds(tmp_path): + staging = tmp_path / ".agentenv-tmp" + staging.mkdir() + hour_ago = (time.time() - 3600, time.time() - 3600) + dead, paused, fresh = staging / "staged-dead", staging / "staged-paused", staging / "staged-fresh" + for f in (dead, paused, fresh, staging / "lock-00"): + f.write_bytes(b"half") + for f in (dead, paused, staging / "lock-00"): + os.utime(f, hour_ago) + holder = _hold_lock(paused) # a live write paused for an hour + try: + LocalFilesystemObjectStore(str(tmp_path)).put("k", b"v") + finally: + holder.kill() + assert not dead.exists() + assert paused.exists() and fresh.exists() and (staging / "lock-00").exists() def test_list_of_an_unwritten_store_is_empty(tmp_path): From 35ea9bc9f60013f9fbb3bc5a43829a06f3ad134b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 30 Sep 2026 10:36:41 -1000 Subject: [PATCH 12/14] fix(store): an overwrite that lands is never reported as failed, and a staged file that cannot be locked is cleaned up Removing the set-aside type record after an overwrite's bytes land, and refreshing its time when it is set aside, are now best effort: a failure there leaves a file the sweep removes, rather than reporting a completed write as failed. A staged file whose lock cannot be taken is closed and removed like any other. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local_object_store.py | 17 +++++++----- tst/unit/store/local_object_store_test.py | 26 +++++++++++++++++++ 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local_object_store.py index 27dbf90c..58038aa1 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local_object_store.py @@ -172,15 +172,16 @@ def _staged(self) -> Iterator[Path]: _sweep(staging) fd, name = tempfile.mkstemp(dir=staging, prefix=_STAGED_PREFIX) staged = Path(name) - if os.name == "nt": # an open file cannot be renamed there, and there is no flock to hold - os.close(fd) - else: - fcntl.flock(fd, fcntl.LOCK_EX) # held while the write lives, so a sweep never takes its file try: + if os.name == "nt": # an open file cannot be renamed there, and there is no flock to hold + os.close(fd) + fd = None + else: + fcntl.flock(fd, fcntl.LOCK_EX) # held while the write lives, so a sweep never takes its file staged.chmod(0o644) yield staged finally: - if os.name != "nt": + if fd is not None: os.close(fd) staged.unlink(missing_ok=True) @@ -207,7 +208,8 @@ def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrit os.replace(aside, meta) raise if aside is not None: - aside.unlink(missing_ok=True) + with contextlib.suppress(OSError): # the bytes have landed; a sweep takes what is left + aside.unlink(missing_ok=True) else: try: os.link(staged, path) # exclusive even against a writer that takes no lock @@ -264,7 +266,8 @@ def _set_aside(self, meta: Path) -> Path | None: os.replace(meta, aside) except FileNotFoundError: return None - os.utime(aside) + with contextlib.suppress(OSError): + os.utime(aside) return aside def _lock_path(self, path: Path) -> Path: diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index 549c00fa..20387986 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -17,6 +17,7 @@ import pytest from agent_env.store import LocalFilesystemObjectStore, ObjectAlreadyExistsError +from agent_env.store.object_store import local_object_store from tst.store import object_conformance @@ -340,6 +341,31 @@ def test_an_overwrite_whose_bytes_do_not_land_keeps_the_old_type(store, tmp_path assert (store.read("cut/x"), store.get_object_metadata("cut/x").content_type) == (b"first", "text/plain") +def test_an_overwrite_whose_cleanup_fails_still_succeeds(store, tmp_path, monkeypatch): + store.put("clean/x", b"first", content_type="text/plain") + real_unlink = Path.unlink + + def unlink(self, missing_ok=False): + if "aside-" in self.name: + raise OSError(errno.EIO, "disk went away") + real_unlink(self, missing_ok=missing_ok) + + monkeypatch.setattr(Path, "unlink", unlink) + store.put("clean/x", b"second!", content_type="application/json", allow_overwrite=True) + assert (store.read("clean/x"), store.get_object_metadata("clean/x").content_type) == (b"second!", "application/json") + + +@_needs_flock +def test_a_write_that_cannot_lock_its_staged_file_leaves_nothing_staged(store, tmp_path, monkeypatch): + def no_lock(fd, op): + raise OSError(errno.ENOLCK, "no locks available") + + monkeypatch.setattr(local_object_store.fcntl, "flock", no_lock) + with pytest.raises(OSError, match="no locks"): + store.put("nolock/x", b"v") + assert [n for n in os.listdir(tmp_path / ".agentenv-tmp") if n.startswith("staged-")] == [] + + def test_an_overwrite_whose_type_cannot_be_recorded_reads_unknown_not_the_old_type(store, tmp_path, monkeypatch): store.put("cut/y", b"first", content_type="text/plain") _failing_replace(monkeypatch, str(tmp_path / ".agentenv-meta" / "cut/y")) From 98f20a3640076ce89840c12f75a5a28b22c73343 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 5 Oct 2026 10:19:58 -0700 Subject: [PATCH 13/14] refactor(store): the local object store and its grant server live together in object_store/local The local backend is now four modules, the store, its grant server, the CA and the grant tokens, and only the local store uses the last three. They move into one package, beside the single-module S3 and GCS backends: local_object_store.py -> local/store.py local_grants/server.py -> local/grant_server.py local_grants/tls.py -> local/tls.py local_grants/tokens.py -> local/tokens.py A rename only: imports and docstrings follow the files, and nothing else changes. LocalFilesystemObjectStore is still exported from agent_env.store.object_store and agent_env.store, which is how everything outside this repository imports it. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/store/object_store/__init__.py | 4 ++-- src/agent_env/store/object_store/local/__init__.py | 6 ++++++ .../{local_grants/server.py => local/grant_server.py} | 6 +++--- .../object_store/{local_object_store.py => local/store.py} | 6 +++--- .../store/object_store/{local_grants => local}/tls.py | 0 .../store/object_store/{local_grants => local}/tokens.py | 0 src/agent_env/store/object_store/local_grants/__init__.py | 2 -- tst/unit/config/test_fixture_prefix_keys.py | 2 +- tst/unit/explorer/explorer_app_test.py | 6 +++--- ...cal_grants_server_test.py => local_grant_server_test.py} | 6 +++--- tst/unit/store/local_object_store_test.py | 4 ++-- .../store/{local_grants_tls_test.py => local_tls_test.py} | 4 ++-- .../{local_grants_tokens_test.py => local_tokens_test.py} | 2 +- tst/unit/store/s3_signed_post_test.py | 2 +- tst/unit/store/test_object_transfer_grants.py | 4 ++-- 15 files changed, 29 insertions(+), 25 deletions(-) create mode 100644 src/agent_env/store/object_store/local/__init__.py rename src/agent_env/store/object_store/{local_grants/server.py => local/grant_server.py} (98%) rename src/agent_env/store/object_store/{local_object_store.py => local/store.py} (98%) rename src/agent_env/store/object_store/{local_grants => local}/tls.py (100%) rename src/agent_env/store/object_store/{local_grants => local}/tokens.py (100%) delete mode 100644 src/agent_env/store/object_store/local_grants/__init__.py rename tst/unit/store/{local_grants_server_test.py => local_grant_server_test.py} (98%) rename tst/unit/store/{local_grants_tls_test.py => local_tls_test.py} (97%) rename tst/unit/store/{local_grants_tokens_test.py => local_tokens_test.py} (95%) diff --git a/src/agent_env/store/object_store/__init__.py b/src/agent_env/store/object_store/__init__.py index 11b82296..414804db 100644 --- a/src/agent_env/store/object_store/__init__.py +++ b/src/agent_env/store/object_store/__init__.py @@ -1,11 +1,11 @@ """Backend-agnostic object store abstraction + its implementations. ``object_store`` defines the abstraction (ObjectStore); ``s3_object_store``, -``local_object_store`` and ``gcs_object_store`` are implementations. The last needs +``local`` and ``gcs_object_store`` are implementations. The last needs the ``gcp`` extra, so it is not re-exported here: an impl pointer names its module. """ -from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore from agent_env.store.object_store.object_store import ( DEFAULT_CONTENT_TYPE, ObjectMetadata, diff --git a/src/agent_env/store/object_store/local/__init__.py b/src/agent_env/store/object_store/local/__init__.py new file mode 100644 index 00000000..b91e507a --- /dev/null +++ b/src/agent_env/store/object_store/local/__init__.py @@ -0,0 +1,6 @@ +"""The local filesystem object store, and the HTTPS transfer grants it issues: signed tokens, a local +certificate authority, and the in-process server that honours them.""" + +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore + +__all__ = ["LocalFilesystemObjectStore"] diff --git a/src/agent_env/store/object_store/local_grants/server.py b/src/agent_env/store/object_store/local/grant_server.py similarity index 98% rename from src/agent_env/store/object_store/local_grants/server.py rename to src/agent_env/store/object_store/local/grant_server.py index 57605fb1..471571fa 100644 --- a/src/agent_env/store/object_store/local_grants/server.py +++ b/src/agent_env/store/object_store/local/grant_server.py @@ -33,8 +33,8 @@ from uvicorn.protocols.http.h11_impl import H11Protocol from agent_env.store.base import GrantUnavailableError -from agent_env.store.object_store.local_grants.tls import check_local_host, local_ca, server_context -from agent_env.store.object_store.local_grants.tokens import ( +from agent_env.store.object_store.local.tls import check_local_host, local_ca, server_context +from agent_env.store.object_store.local.tokens import ( GrantClaims, GrantSigner, InvalidGrantError, @@ -44,7 +44,7 @@ from agent_env.store.object_store.object_store import DEFAULT_CONTENT_TYPE if TYPE_CHECKING: - from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore + from agent_env.store.object_store.local.store import LocalFilesystemObjectStore logger = logging.getLogger(__name__) diff --git a/src/agent_env/store/object_store/local_object_store.py b/src/agent_env/store/object_store/local/store.py similarity index 98% rename from src/agent_env/store/object_store/local_object_store.py rename to src/agent_env/store/object_store/local/store.py index 1439b35c..acbef743 100644 --- a/src/agent_env/store/object_store/local_object_store.py +++ b/src/agent_env/store/object_store/local/store.py @@ -26,9 +26,9 @@ if os.name != "nt": import fcntl from agent_env.store.local_state import ensure_state_dir -from agent_env.store.object_store.local_grants.server import grant_server -from agent_env.store.object_store.local_grants.tls import check_local_host -from agent_env.store.object_store.local_grants.tokens import Op +from agent_env.store.object_store.local.grant_server import grant_server +from agent_env.store.object_store.local.tls import check_local_host +from agent_env.store.object_store.local.tokens import Op from agent_env.store.object_store.object_store import ( DEFAULT_CONTENT_TYPE, ObjectMetadata, diff --git a/src/agent_env/store/object_store/local_grants/tls.py b/src/agent_env/store/object_store/local/tls.py similarity index 100% rename from src/agent_env/store/object_store/local_grants/tls.py rename to src/agent_env/store/object_store/local/tls.py diff --git a/src/agent_env/store/object_store/local_grants/tokens.py b/src/agent_env/store/object_store/local/tokens.py similarity index 100% rename from src/agent_env/store/object_store/local_grants/tokens.py rename to src/agent_env/store/object_store/local/tokens.py diff --git a/src/agent_env/store/object_store/local_grants/__init__.py b/src/agent_env/store/object_store/local_grants/__init__.py deleted file mode 100644 index 4fdd242f..00000000 --- a/src/agent_env/store/object_store/local_grants/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -"""HTTPS object-transfer grants for the local filesystem object store: signed tokens, a local -certificate authority, and the in-process server that honours them.""" diff --git a/tst/unit/config/test_fixture_prefix_keys.py b/tst/unit/config/test_fixture_prefix_keys.py index e4de0ec0..0b19a3e8 100644 --- a/tst/unit/config/test_fixture_prefix_keys.py +++ b/tst/unit/config/test_fixture_prefix_keys.py @@ -12,7 +12,7 @@ from agent_env.artifact.store import reset_artifact_store from agent_env.config import configure, get_config, set_object_store from agent_env.env.snapshot_store import _snapshot_servicedb -from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep from agent_env.task_step.task_steps.verifiers.judge_utils.trajectory_filter import TrajectoryFilter from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep diff --git a/tst/unit/explorer/explorer_app_test.py b/tst/unit/explorer/explorer_app_test.py index 1925f02b..8437518b 100644 --- a/tst/unit/explorer/explorer_app_test.py +++ b/tst/unit/explorer/explorer_app_test.py @@ -19,7 +19,7 @@ from agent_env.config.errors import ConfigError from agent_env.runner import store as run_store from agent_env.runner.local_runner import LocalRunner -from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore from agent_env.store.document_store.sqlite_document_store import LocalSqliteDocumentStore from agent_env.task import Task from agent_env.explorer.routers import objects as objects_router @@ -591,7 +591,7 @@ def test_root_serves_health_when_no_ui_is_mounted(client): def test_instance_content_serves_object_store_bytes(client, tmp_path): """Content route proxies an object-store blob and refuses a url outside the store.""" from agent_env.config import set_object_store - from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore + from agent_env.store.object_store.local.store import LocalFilesystemObjectStore store = LocalFilesystemObjectStore(str(tmp_path / "obj")) set_object_store(store) @@ -698,7 +698,7 @@ def test_content_sandboxes_active_types(client, tmp_path): active type and add a misleading policy to a passive one. """ from agent_env.config import set_object_store - from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore + from agent_env.store.object_store.local.store import LocalFilesystemObjectStore store = LocalFilesystemObjectStore(str(tmp_path / "obj")) set_object_store(store) diff --git a/tst/unit/store/local_grants_server_test.py b/tst/unit/store/local_grant_server_test.py similarity index 98% rename from tst/unit/store/local_grants_server_test.py rename to tst/unit/store/local_grant_server_test.py index 96ebdb38..9e74f1b8 100644 --- a/tst/unit/store/local_grants_server_test.py +++ b/tst/unit/store/local_grant_server_test.py @@ -13,9 +13,9 @@ from agent_env.config.paths import state_root from agent_env.store import LocalFilesystemObjectStore -from agent_env.store.object_store.local_grants import server as server_module -from agent_env.store.object_store.local_grants.server import default_bind_host, grant_server -from agent_env.store.object_store.local_grants.tls import local_ca +from agent_env.store.object_store.local import grant_server as server_module +from agent_env.store.object_store.local.grant_server import default_bind_host, grant_server +from agent_env.store.object_store.local.tls import local_ca from tst.store import object_conformance diff --git a/tst/unit/store/local_object_store_test.py b/tst/unit/store/local_object_store_test.py index 20387986..a79ce5a8 100644 --- a/tst/unit/store/local_object_store_test.py +++ b/tst/unit/store/local_object_store_test.py @@ -17,7 +17,7 @@ import pytest from agent_env.store import LocalFilesystemObjectStore, ObjectAlreadyExistsError -from agent_env.store.object_store import local_object_store +from agent_env.store.object_store.local import store as local_store from tst.store import object_conformance @@ -360,7 +360,7 @@ def test_a_write_that_cannot_lock_its_staged_file_leaves_nothing_staged(store, t def no_lock(fd, op): raise OSError(errno.ENOLCK, "no locks available") - monkeypatch.setattr(local_object_store.fcntl, "flock", no_lock) + monkeypatch.setattr(local_store.fcntl, "flock", no_lock) with pytest.raises(OSError, match="no locks"): store.put("nolock/x", b"v") assert [n for n in os.listdir(tmp_path / ".agentenv-tmp") if n.startswith("staged-")] == [] diff --git a/tst/unit/store/local_grants_tls_test.py b/tst/unit/store/local_tls_test.py similarity index 97% rename from tst/unit/store/local_grants_tls_test.py rename to tst/unit/store/local_tls_test.py index a1282984..9c557472 100644 --- a/tst/unit/store/local_grants_tls_test.py +++ b/tst/unit/store/local_tls_test.py @@ -12,8 +12,8 @@ from cryptography import x509 from agent_env.config.paths import state_root -from agent_env.store.object_store.local_grants import tls -from agent_env.store.object_store.local_grants.tls import check_local_host, local_ca, server_context +from agent_env.store.object_store.local import tls +from agent_env.store.object_store.local.tls import check_local_host, local_ca, server_context def _handshake(server: ssl.SSLContext, *, cafile: Path, hostname: str) -> None: diff --git a/tst/unit/store/local_grants_tokens_test.py b/tst/unit/store/local_tokens_test.py similarity index 95% rename from tst/unit/store/local_grants_tokens_test.py rename to tst/unit/store/local_tokens_test.py index 9ed7d427..bc07880d 100644 --- a/tst/unit/store/local_grants_tokens_test.py +++ b/tst/unit/store/local_tokens_test.py @@ -2,7 +2,7 @@ import pytest -from agent_env.store.object_store.local_grants.tokens import GrantClaims, GrantSigner, InvalidGrantError +from agent_env.store.object_store.local.tokens import GrantClaims, GrantSigner, InvalidGrantError NOW = 1_800_000_000 diff --git a/tst/unit/store/s3_signed_post_test.py b/tst/unit/store/s3_signed_post_test.py index bdf3c389..93140d4c 100644 --- a/tst/unit/store/s3_signed_post_test.py +++ b/tst/unit/store/s3_signed_post_test.py @@ -2,7 +2,7 @@ it — the round trip is integration territory.""" from __future__ import annotations -from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore from agent_env.store.object_store.s3_object_store import S3ObjectStore BUCKET = "artifact-bucket" diff --git a/tst/unit/store/test_object_transfer_grants.py b/tst/unit/store/test_object_transfer_grants.py index 9027dd56..dda7ea9d 100644 --- a/tst/unit/store/test_object_transfer_grants.py +++ b/tst/unit/store/test_object_transfer_grants.py @@ -16,7 +16,7 @@ from botocore.credentials import Credentials, RefreshableCredentials from agent_env.store import GrantUnavailableError -from agent_env.store.object_store.local_object_store import ( +from agent_env.store.object_store.local.store import ( LocalFilesystemObjectStore, ) from agent_env.store.object_store import UploadPolicy @@ -282,7 +282,7 @@ def test_botocore_still_exposes_the_signing_credentials() -> None: def test_local_store_does_not_offer_transfer_grants(tmp_path) -> None: - """It can issue them (local_grants_server_test.py), but agents on the local store don't use them.""" + """It can issue them (local_grant_server_test.py), but agents on the local store don't use them.""" store = LocalFilesystemObjectStore(str(tmp_path)) assert S3ObjectStore.supports_transfer_grants From 9ba5dea626fca6797af00a3d1e2c7a480200d065 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 5 Oct 2026 11:36:14 -0700 Subject: [PATCH 14/14] fix(store): a HEAD through a local read grant reports the size and type without reading the object Starlette's StreamingResponse has no HEAD case, so a HEAD streamed the whole object only for the server to drop the body. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../store/object_store/local/grant_server.py | 12 +++++++----- tst/unit/store/local_grant_server_test.py | 3 ++- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/agent_env/store/object_store/local/grant_server.py b/src/agent_env/store/object_store/local/grant_server.py index 471571fa..89725141 100644 --- a/src/agent_env/store/object_store/local/grant_server.py +++ b/src/agent_env/store/object_store/local/grant_server.py @@ -217,7 +217,7 @@ async def _handle(self, request: Request) -> Response: return PlainTextResponse("The grant does not allow this request.", status_code=403) try: if claims.op == "get": - return _get(store, claims) + return _get(store, claims, head=request.method == "HEAD") if claims.op == "put": await _put(store, claims, request) return Response(status_code=200) @@ -239,7 +239,7 @@ def _path(store: LocalFilesystemObjectStore, key: str) -> Path: raise _Rejected(403, "The grant's key is not an object key of this store.") from None -def _get(store: LocalFilesystemObjectStore, claims: GrantClaims) -> Response: +def _get(store: LocalFilesystemObjectStore, claims: GrantClaims, *, head: bool = False) -> Response: path = _path(store, claims.key) # Opened under the key's lock and typed from that open file, so the bytes sent and their type are one write's. with store._locked(path, shared=True): @@ -249,9 +249,11 @@ def _get(store: LocalFilesystemObjectStore, claims: GrantClaims) -> Response: raise _Rejected(404, "No object exists at this grant's key.") from None st = os.fstat(f.fileno()) content_type = store._read_content_type(path, st) - return StreamingResponse( - _chunks(f), media_type=content_type or DEFAULT_CONTENT_TYPE, headers={"Content-Length": str(st.st_size)} - ) + headers = {"Content-Length": str(st.st_size)} + if head: # the size and type alone: a StreamingResponse would read the whole object to send nothing + f.close() + return Response(media_type=content_type or DEFAULT_CONTENT_TYPE, headers=headers) + return StreamingResponse(_chunks(f), media_type=content_type or DEFAULT_CONTENT_TYPE, headers=headers) async def _chunks(f: BinaryIO) -> AsyncIterator[bytes]: diff --git a/tst/unit/store/local_grant_server_test.py b/tst/unit/store/local_grant_server_test.py index 9e74f1b8..a804b803 100644 --- a/tst/unit/store/local_grant_server_test.py +++ b/tst/unit/store/local_grant_server_test.py @@ -99,7 +99,8 @@ def test_a_read_of_a_missing_object_is_404(store): assert httpx.get(store.issue_read_grant(store.object_url("a/none.bin")).url).status_code == 404 -def test_a_head_reports_the_size(store): +def test_a_head_reports_the_size_without_reading_the_object(store, monkeypatch): + monkeypatch.setattr(server_module, "_chunks", lambda f: pytest.fail("a HEAD read the object")) url = store.put("a/b.bin", b"12345") response = httpx.head(store.issue_read_grant(url).url) assert response.status_code == 200