diff --git a/src/agent_env/a2a_agent/object_transfer.py b/src/agent_env/a2a_agent/object_transfer.py index 71471c35..4736eff8 100644 --- a/src/agent_env/a2a_agent/object_transfer.py +++ b/src/agent_env/a2a_agent/object_transfer.py @@ -50,10 +50,10 @@ _INCREMENT_NAME = re.compile(r"^(?P[0-9]{6})(?:\.[A-Za-z0-9][A-Za-z0-9._-]*)?$") _TRANSFER_UNAVAILABLE = "transfer_unavailable" # the SDK's code for a store it could not reach -# The time budget of one transfer, outermost first: a grant outlives agent-env's wait for the -# agent's answer, which outlasts the SDK's retries of stalled connections +# The time budget of one transfer, outermost first: a grant (the issuing store's +# grant_lifetime_seconds, 12 hours unless configured) outlives agent-env's wait for the agent's +# answer, which outlasts the SDK's retries of stalled connections # (agentenv_protocol.transfers.TRANSFER_STALL_BUDGET_SECONDS). -GRANT_LIFETIME_SECONDS = 3600 TRANSFER_TIMEOUT_SECONDS = 600 # a call during which the agent moves objects REPLY_TIMEOUT_SECONDS = 120 # a call during which it moves none @@ -198,9 +198,7 @@ def write_object(store: ObjectStore, url: str, *, media_type: str, max_bytes: in return WriteObject( media_type=media_type, max_bytes=max_bytes, - write=store.issue_write_grant( - url, media_type=media_type, max_bytes=max_bytes, expires_in=GRANT_LIFETIME_SECONDS - ), + write=store.issue_write_grant(url, media_type=media_type, max_bytes=max_bytes), ) @@ -219,7 +217,7 @@ def read_object( media_type=media_type or stored_type, max_bytes=max_bytes or max(metadata.size, 1), size_bytes=metadata.size, - read=store.issue_read_grant(url, expires_in=GRANT_LIFETIME_SECONDS), + read=store.issue_read_grant(url), ) diff --git a/src/agent_env/store/object_store/__init__.py b/src/agent_env/store/object_store/__init__.py index 414804db..643f2e92 100644 --- a/src/agent_env/store/object_store/__init__.py +++ b/src/agent_env/store/object_store/__init__.py @@ -8,6 +8,8 @@ from agent_env.store.object_store.local.store import LocalFilesystemObjectStore from agent_env.store.object_store.object_store import ( DEFAULT_CONTENT_TYPE, + DEFAULT_GRANT_LIFETIME_SECONDS, + MIN_GRANT_LIFETIME_SECONDS, ObjectMetadata, ObjectStore, UploadPolicy, @@ -19,6 +21,8 @@ "ObjectMetadata", "UploadPolicy", "DEFAULT_CONTENT_TYPE", + "DEFAULT_GRANT_LIFETIME_SECONDS", + "MIN_GRANT_LIFETIME_SECONDS", "S3ObjectStore", "LocalFilesystemObjectStore", ] diff --git a/src/agent_env/store/object_store/gcs_object_store.py b/src/agent_env/store/object_store/gcs_object_store.py index 52ebaa5a..e8c5d2f7 100644 --- a/src/agent_env/store/object_store/gcs_object_store.py +++ b/src/agent_env/store/object_store/gcs_object_store.py @@ -31,9 +31,11 @@ from agent_env.store.base import GrantUnavailableError, ObjectAlreadyExistsError, ObjectNotFoundError from agent_env.store.object_store.object_store import ( DEFAULT_CONTENT_TYPE, + DEFAULT_GRANT_LIFETIME_SECONDS, ObjectMetadata, ObjectStore, UploadPolicy, + grant_lifetime, ) logger = logging.getLogger(__name__) @@ -84,7 +86,14 @@ class GcsObjectStore(ObjectStore): max_single_upload_bytes = 5 * 1024**4 # one Cloud Storage object - def __init__(self, client, bucket: str, *, signer: Signing | None = None) -> None: + def __init__( + self, + client, + bucket: str, + *, + signer: Signing | None = None, + grant_lifetime_seconds: int = DEFAULT_GRANT_LIFETIME_SECONDS, + ) -> None: self._client = client self._bucket = bucket self._signer = signer @@ -93,11 +102,17 @@ def __init__(self, client, bucket: str, *, signer: Signing | None = None) -> Non self._max_signed_seconds = ( _MAX_KEY_SIGNED_SECONDS if isinstance(signer, service_account.Credentials) else _MAX_IAM_SIGNED_SECONDS ) + self.grant_lifetime_seconds = grant_lifetime(grant_lifetime_seconds, most=self._max_signed_seconds) self._warned_unsigned = False @classmethod def from_config( - cls, *, bucket: str, project: str | None = None, signing_service_account: str | None = None + cls, + *, + bucket: str, + project: str | None = None, + signing_service_account: str | None = None, + grant_lifetime_seconds: int = DEFAULT_GRANT_LIFETIME_SECONDS, ) -> GcsObjectStore: """Authenticate through Application Default Credentials. With ``signing_service_account`` the store always signs as that account through IAM, which needs @@ -127,7 +142,7 @@ def from_config( quota = getattr(credentials, "quota_project_id", None) billed = f" (requests are billed to quota project {quota!r})" if quota else "" raise ConfigError(f"Cannot sign as {signer.signer_email}{billed}: {e}") from e - return cls(client, bucket, signer=signer) + return cls(client, bucket, signer=signer, grant_lifetime_seconds=grant_lifetime_seconds) def put( self, @@ -230,7 +245,8 @@ def signed_put_url(self, object_url: str, expires_in: int = 3600) -> str | None: return None return self._sign(blob, "PUT", min(expires_in, self._max_signed_seconds)) - def issue_read_grant(self, object_url: str, *, expires_in: int = 3600) -> HttpGetGrant: + def issue_read_grant(self, object_url: str, *, expires_in: int | None = None) -> HttpGetGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in blob = self._blob(object_url) expires_at = self._grant_expiry(expires_in) return HttpGetGrant(kind="http-get", url=self._sign(blob, "GET", expires_in), expires_at=expires_at) @@ -241,8 +257,9 @@ def issue_write_grant( *, media_type: str, max_bytes: int, - expires_in: int = 3600, + expires_in: int | None = None, ) -> HttpPutGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in blob = self._blob(object_url) expires_at = self._grant_expiry(expires_in) bound = {"x-goog-content-length-range": f"0,{max_bytes}"} diff --git a/src/agent_env/store/object_store/local/store.py b/src/agent_env/store/object_store/local/store.py index 296a41bb..932c0a9e 100644 --- a/src/agent_env/store/object_store/local/store.py +++ b/src/agent_env/store/object_store/local/store.py @@ -31,9 +31,11 @@ from agent_env.store.object_store.local.tokens import Op from agent_env.store.object_store.object_store import ( DEFAULT_CONTENT_TYPE, + DEFAULT_GRANT_LIFETIME_SECONDS, ObjectMetadata, ObjectStore, UploadPolicy, + grant_lifetime, ) logger = logging.getLogger(__name__) @@ -47,6 +49,7 @@ _STAGED_PREFIX = "staged-" _STAGED_GRACE_SECONDS = 60 # a staged file younger than this may not be locked by its writer yet _GRANT_MODES = ("auto", "off") +_MAX_GRANT_LIFETIME_SECONDS = 7 * 24 * 60 * 60 # as S3's; a grant cannot outlive its process anyway _LOCAL_SANDBOX_TYPE = "local" # LocalSandbox.type: containers on this host's Docker, which trust the local CA logger = logging.getLogger(__name__) @@ -80,6 +83,7 @@ def __init__( grants: str = "auto", grant_bind_host: str | None = None, grant_advertise_host: str | None = None, + grant_lifetime_seconds: int = DEFAULT_GRANT_LIFETIME_SECONDS, ) -> None: if grants not in _GRANT_MODES: raise ValueError(f"grants must be one of {_GRANT_MODES}, got {grants!r}") @@ -92,6 +96,7 @@ def __init__( check_local_host(grant_advertise_host) self._root = Path(root) if root is not None else state_root() / "object_store" self.supports_transfer_grants = grants == "auto" + self.grant_lifetime_seconds = grant_lifetime(grant_lifetime_seconds, most=_MAX_GRANT_LIFETIME_SECONDS) self._grant_bind_host = grant_bind_host self._grant_advertise_host = grant_advertise_host self._swept = False @@ -191,13 +196,15 @@ def get_object_key(self, object_url: str) -> str: def grants_reach(self, sandbox_type: str | None) -> bool: return sandbox_type == _LOCAL_SANDBOX_TYPE - def issue_read_grant(self, object_url: str, *, expires_in: int = 3600) -> HttpGetGrant: + def issue_read_grant(self, object_url: str, *, expires_in: int | None = None) -> HttpGetGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in 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 + self, object_url: str, *, media_type: str, max_bytes: int, expires_in: int | None = None ) -> HttpPutGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in url, expires_at = self._grant( "put", expires_in, key=self._grant_key(object_url), max_bytes=max_bytes, content_type=media_type ) diff --git a/src/agent_env/store/object_store/object_store.py b/src/agent_env/store/object_store/object_store.py index ec0fa2da..940a3bb2 100644 --- a/src/agent_env/store/object_store/object_store.py +++ b/src/agent_env/store/object_store/object_store.py @@ -15,6 +15,24 @@ ) DEFAULT_CONTENT_TYPE = "application/octet-stream" +DEFAULT_GRANT_LIFETIME_SECONDS = 12 * 60 * 60 +# The shortest a store's grants may be set to last: longer than the ten minutes agent-env waits for an agent +# to move an object through one. +MIN_GRANT_LIFETIME_SECONDS = 15 * 60 + + +def grant_lifetime(seconds: object, *, most: int | None = None) -> int: + """``seconds`` checked as a store's ``grant_lifetime_seconds`` setting: a whole number of seconds, at + least ``MIN_GRANT_LIFETIME_SECONDS``, and no more than ``most`` where the store cannot sign for longer.""" + if isinstance(seconds, bool) or not isinstance(seconds, int): + raise ValueError(f"grant_lifetime_seconds must be a whole number of seconds, got {seconds!r}") + if seconds < MIN_GRANT_LIFETIME_SECONDS: + raise ValueError( + f"grant_lifetime_seconds must be at least {MIN_GRANT_LIFETIME_SECONDS}, to outlast a transfer; got {seconds}" + ) + if most is not None and seconds > most: + raise ValueError(f"grant_lifetime_seconds can be at most {most} for this store, got {seconds}") + return seconds @dataclass(frozen=True) @@ -47,6 +65,9 @@ class ObjectStore(ABC): supports_transfer_grants: bool = False # The largest object one upload through a grant can create, where the provider caps it. max_single_upload_bytes: int | None = None + # How long a read or write grant lasts when its caller names no lifetime; the built-in stores + # take it as their grant_lifetime_seconds setting. + grant_lifetime_seconds: int = DEFAULT_GRANT_LIFETIME_SECONDS @classmethod def from_config(cls, **config) -> ObjectStore: @@ -167,11 +188,12 @@ def grants_reach(self, sandbox_type: str | None) -> bool: return True def issue_read_grant( - self, object_url: str, *, expires_in: int = 3600 + self, object_url: str, *, expires_in: int | None = None ) -> HttpGetGrant: """An HTTPS GET grant for the existing object at ``object_url``. ``expires_at`` is at most - ``expires_in`` away, and earlier if the store's signing credentials expire first; raise - GrantUnavailableError for an ``expires_in`` beyond what the store can sign.""" + ``expires_in`` away (None: ``grant_lifetime_seconds``), and earlier if the store's signing + credentials expire first; raise GrantUnavailableError for an ``expires_in`` beyond what the + store can sign.""" raise NotImplementedError( f"{type(self).__name__} cannot issue remote object-transfer grants" ) @@ -182,7 +204,7 @@ def issue_write_grant( *, media_type: str, max_bytes: int, - expires_in: int = 3600, + expires_in: int | None = None, ) -> HttpPutGrant: """An HTTPS PUT grant for one object at ``object_url``, signed for ``media_type`` and bounded to ``max_bytes`` where the provider can enforce that; expiry as for a read grant.""" diff --git a/src/agent_env/store/object_store/s3_object_store.py b/src/agent_env/store/object_store/s3_object_store.py index 17629ebc..259d7eca 100644 --- a/src/agent_env/store/object_store/s3_object_store.py +++ b/src/agent_env/store/object_store/s3_object_store.py @@ -22,9 +22,11 @@ from agent_env.store.base import GrantUnavailableError, ObjectAlreadyExistsError, ObjectNotFoundError from agent_env.store.object_store.object_store import ( DEFAULT_CONTENT_TYPE, + DEFAULT_GRANT_LIFETIME_SECONDS, ObjectMetadata, ObjectStore, UploadPolicy, + grant_lifetime, ) logger = logging.getLogger(__name__) @@ -51,10 +53,17 @@ class S3ObjectStore(ObjectStore): max_single_upload_bytes = 5 * 1024 * 1024 * 1024 # one S3 PUT def __init__( - self, client, bucket: str, *, region: str | None = None, share_credentials: bool = False + self, + client, + bucket: str, + *, + region: str | None = None, + share_credentials: bool = False, + grant_lifetime_seconds: int = DEFAULT_GRANT_LIFETIME_SECONDS, ) -> None: if not isinstance(share_credentials, bool): raise ValueError(f"share_credentials must be true or false, got {share_credentials!r}") + self.grant_lifetime_seconds = grant_lifetime(grant_lifetime_seconds, most=_MAX_SIGV4_EXPIRY_SECONDS) self._s3 = client self._bucket = bucket self._region = region @@ -66,7 +75,12 @@ def __init__( @classmethod def from_config( - cls, *, bucket: str, region: str | None = None, share_credentials: bool = False + cls, + *, + bucket: str, + region: str | None = None, + share_credentials: bool = False, + grant_lifetime_seconds: int = DEFAULT_GRANT_LIFETIME_SECONDS, ) -> S3ObjectStore: """Build an adaptive-retry boto3 S3 client for ``bucket``.""" kwargs = { @@ -78,7 +92,10 @@ def from_config( } if region: kwargs["region_name"] = region - return cls(boto3.client("s3", **kwargs), bucket, region=region, share_credentials=share_credentials) + return cls( + boto3.client("s3", **kwargs), bucket, region=region, share_credentials=share_credentials, + grant_lifetime_seconds=grant_lifetime_seconds, + ) @property def bucket(self) -> str: @@ -164,8 +181,9 @@ def signed_put_url(self, object_url: str, expires_in: int = 3600) -> str | None: return self._signed_url("put_object", object_url, expires_in) def issue_read_grant( - self, object_url: str, *, expires_in: int = 3600 + self, object_url: str, *, expires_in: int | None = None ) -> HttpGetGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in return HttpGetGrant( kind="http-get", url=self._signed_url("get_object", object_url, expires_in), @@ -178,8 +196,9 @@ def issue_write_grant( *, media_type: str, max_bytes: int, - expires_in: int = 3600, + expires_in: int | None = None, ) -> HttpPutGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in # A presigned S3 PUT cannot bound the upload size, so the uploader enforces max_bytes. return HttpPutGrant( kind="http-put", diff --git a/src/agent_env/store/routing.py b/src/agent_env/store/routing.py index 4fe2bd20..282c5622 100644 --- a/src/agent_env/store/routing.py +++ b/src/agent_env/store/routing.py @@ -558,12 +558,14 @@ def signed_post(self, url_prefix: str, *, expires_in: int = 3600, max_bytes: int def shared_credentials_env(self) -> dict[str, str]: return self.configured.shared_credentials_env() - def issue_read_grant(self, object_url: str, *, expires_in: int = 3600) -> HttpGetGrant: - return self._at(object_url).issue_read_grant(object_url, expires_in=expires_in) + def issue_read_grant(self, object_url: str, *, expires_in: int | None = None) -> HttpGetGrant: + return self._at(object_url).issue_read_grant(object_url, **_lifetime(expires_in)) - def issue_write_grant(self, object_url: str, *, media_type: str, max_bytes: int, expires_in: int = 3600) -> HttpPutGrant: + def issue_write_grant( + self, object_url: str, *, media_type: str, max_bytes: int, expires_in: int | None = None + ) -> HttpPutGrant: return self._writing(object_url).issue_write_grant( - object_url, media_type=media_type, max_bytes=max_bytes, expires_in=expires_in + object_url, media_type=media_type, max_bytes=max_bytes, **_lifetime(expires_in) ) def issue_upload_policy(self, prefix_url: str, *, max_object_bytes: int, expires_in: int) -> UploadPolicy: @@ -604,3 +606,9 @@ def auth(self, ref: str) -> RegistryAuth | None: def _for(self, repository: str) -> ImageStore: return self.local if repository.startswith(_LOCAL_REPOSITORY_PREFIX) else self.configured + + +def _lifetime(expires_in: int | None) -> dict[str, int]: + """``expires_in`` as a keyword only when the caller named one, so the owning store applies its own default, + whatever it is.""" + return {} if expires_in is None else {"expires_in": expires_in} diff --git a/tst/unit/a2a_agent/object_transfer_test.py b/tst/unit/a2a_agent/object_transfer_test.py index 1cdac0ce..84619ddf 100644 --- a/tst/unit/a2a_agent/object_transfer_test.py +++ b/tst/unit/a2a_agent/object_transfer_test.py @@ -19,6 +19,7 @@ read_objects_under, write_object, ) +from agent_env.store.object_store import DEFAULT_GRANT_LIFETIME_SECONDS from tst.util.granting_object_store import GrantingObjectStore @@ -253,7 +254,7 @@ def test_the_transfer_time_budget_nests(): assert ( TRANSFER_STALL_BUDGET_SECONDS < object_transfer.TRANSFER_TIMEOUT_SECONDS - < object_transfer.GRANT_LIFETIME_SECONDS + < DEFAULT_GRANT_LIFETIME_SECONDS ) diff --git a/tst/unit/a2a_agent/validator_trajectory_dag_test.py b/tst/unit/a2a_agent/validator_trajectory_dag_test.py index 324bc71c..cc58b82b 100644 --- a/tst/unit/a2a_agent/validator_trajectory_dag_test.py +++ b/tst/unit/a2a_agent/validator_trajectory_dag_test.py @@ -227,7 +227,7 @@ def grants_reach(self, sandbox_type: str | None) -> bool: def get_object_key(self, object_url: str) -> str: return object_url.removeprefix("s3://test-bucket/") - def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in): + def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in=None): self.write_grants.append(object_url) return HttpPutGrant( kind="http-put", diff --git a/tst/unit/store/gcs_object_store_test.py b/tst/unit/store/gcs_object_store_test.py index dd0eeb4a..c9304075 100644 --- a/tst/unit/store/gcs_object_store_test.py +++ b/tst/unit/store/gcs_object_store_test.py @@ -26,7 +26,7 @@ from agent_env.a2a_agent.object_transfer import changelog_enable_call from agent_env.config.errors import ConfigError from agent_env.store import GrantUnavailableError, ObjectAlreadyExistsError, ObjectNotFoundError, _google -from agent_env.store.object_store import gcs_object_store +from agent_env.store.object_store import DEFAULT_GRANT_LIFETIME_SECONDS, gcs_object_store from agent_env.store.object_store.gcs_object_store import GcsObjectStore from tst.store import object_conformance @@ -408,6 +408,19 @@ def test_urls_signed_with_a_key_last_up_to_seven_days(client): store.issue_read_grant(f"gs://{BUCKET}/k", expires_in=604801) +def test_grants_last_the_configured_lifetime(client, signer): + store = GcsObjectStore(client, BUCKET, signer=signer, grant_lifetime_seconds=900) + _assert_expiry(store.issue_read_grant(store.put("k", b"v")).expires_at, seconds=900) + + +@pytest.mark.parametrize("signing, most", [("iam", 43200), ("key", 604800)]) +def test_a_lifetime_the_signer_cannot_sign_is_refused_when_the_store_is_made(client, signing, most): + signer = _FakeSigner() if signing == "iam" else _key_credentials() + GcsObjectStore(client, BUCKET, signer=signer, grant_lifetime_seconds=most) + with pytest.raises(ValueError, match=f"at most {most}"): + GcsObjectStore(client, BUCKET, signer=signer, grant_lifetime_seconds=most + 1) + + def test_without_a_signer_signed_urls_are_absent_and_warn_once(unsigned, caplog): with caplog.at_level(logging.WARNING, logger=gcs_object_store.__name__): assert unsigned.signed_get_url(f"gs://{BUCKET}/k") is None @@ -444,7 +457,7 @@ def test_write_grant_signs_its_media_type_and_size_bound(store): assert "content-type=application/json" in grant.url assert "x-goog-content-length-range" in grant.url assert grant.headers == {"Content-Type": "application/json", "x-goog-content-length-range": "0,64"} - _assert_expiry(grant.expires_at, seconds=3600) + _assert_expiry(grant.expires_at, seconds=DEFAULT_GRANT_LIFETIME_SECONDS) def _policy(fields: dict) -> dict: diff --git a/tst/unit/store/grant_lifetime_test.py b/tst/unit/store/grant_lifetime_test.py new file mode 100644 index 00000000..45ea7d43 --- /dev/null +++ b/tst/unit/store/grant_lifetime_test.py @@ -0,0 +1,110 @@ +"""Each store sets how long its read and write grants last, 12 hours unless configured; agent-env names no +lifetime of its own, so a grant lasts what the store that issues it says.""" + +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace + +import pytest + +from agent_env.a2a_agent.object_transfer import TRANSFER_TIMEOUT_SECONDS, read_object, write_object +from agent_env.config import Config +from agent_env.store import LocalFilesystemObjectStore +from agent_env.store.object_store import DEFAULT_GRANT_LIFETIME_SECONDS, MIN_GRANT_LIFETIME_SECONDS, S3ObjectStore +from agent_env.store.routing import LocalRunObjectStore +from tst.util.granting_object_store import GrantingObjectStore + +BUCKET = "artifact-bucket" + + +class _StubS3: + def __init__(self) -> None: + self._request_signer = SimpleNamespace(_credentials=SimpleNamespace(token=None)) # long-term + + def generate_presigned_url(self, operation, **kwargs): + return f"https://{BUCKET}.s3.amazonaws.com/signed" + + +def _lasts(expires_at: datetime, seconds: int) -> bool: + return abs((expires_at - (datetime.now(UTC) + timedelta(seconds=seconds))).total_seconds()) < 2 + + +def test_the_default_is_twelve_hours(): + assert DEFAULT_GRANT_LIFETIME_SECONDS == 12 * 60 * 60 + store = S3ObjectStore(_StubS3(), BUCKET) + assert _lasts(store.issue_read_grant(f"s3://{BUCKET}/k").expires_at, DEFAULT_GRANT_LIFETIME_SECONDS) + + +def test_an_s3_store_grants_for_its_configured_lifetime(): + store = S3ObjectStore(_StubS3(), BUCKET, grant_lifetime_seconds=900) + assert _lasts(store.issue_read_grant(f"s3://{BUCKET}/k").expires_at, 900) + write = store.issue_write_grant(f"s3://{BUCKET}/k", media_type="application/json", max_bytes=8) + assert _lasts(write.expires_at, 900) + + +def test_a_caller_can_still_name_a_lifetime(): + store = S3ObjectStore(_StubS3(), BUCKET, grant_lifetime_seconds=900) + assert _lasts(store.issue_read_grant(f"s3://{BUCKET}/k", expires_in=60).expires_at, 60) + + +@pytest.mark.parametrize("seconds, why", [ + (True, "whole number"), (3600.0, "whole number"), ("3600", "whole number"), + (0, "at least 900"), (-1, "at least 900"), (899, "at least 900"), +]) +def test_a_lifetime_that_is_not_whole_seconds_or_is_too_short_to_outlast_a_transfer_is_refused(seconds, why, tmp_path): + with pytest.raises(ValueError, match=why): + S3ObjectStore(_StubS3(), BUCKET, grant_lifetime_seconds=seconds) + with pytest.raises(ValueError, match=why): + LocalFilesystemObjectStore(str(tmp_path), grant_lifetime_seconds=seconds) + + +def test_the_shortest_lifetime_outlasts_agent_envs_wait_for_a_transfer(): + assert TRANSFER_TIMEOUT_SECONDS < MIN_GRANT_LIFETIME_SECONDS <= DEFAULT_GRANT_LIFETIME_SECONDS + + +def test_a_local_lifetime_beyond_seven_days_is_refused(tmp_path): + LocalFilesystemObjectStore(str(tmp_path), grant_lifetime_seconds=7 * 24 * 60 * 60) + with pytest.raises(ValueError, match="at most 604800"): + LocalFilesystemObjectStore(str(tmp_path), grant_lifetime_seconds=7 * 24 * 60 * 60 + 1) + + +def test_an_s3_lifetime_beyond_sigv4s_seven_days_is_refused(): + S3ObjectStore(_StubS3(), BUCKET, grant_lifetime_seconds=7 * 24 * 60 * 60) + with pytest.raises(ValueError, match="at most 604800"): + S3ObjectStore(_StubS3(), BUCKET, grant_lifetime_seconds=7 * 24 * 60 * 60 + 1) + + +def test_agent_env_grants_for_the_issuing_stores_lifetime(tmp_path): + store = GrantingObjectStore(str(tmp_path)) + store.grant_lifetime_seconds = 900 + url = store.put("a/in.json", b"{}", content_type="application/json") + assert _lasts(read_object(store, url).read.expires_at, 900) + assert _lasts(write_object(store, store.object_url("a/out.json"), media_type="application/json", max_bytes=8).write.expires_at, 900) + + +def test_a_local_run_grants_for_the_lifetime_of_the_store_that_owns_the_object(tmp_path): + local = GrantingObjectStore(str(tmp_path)) + local.grant_lifetime_seconds = 1800 + routed = LocalRunObjectStore(S3ObjectStore(_StubS3(), BUCKET, grant_lifetime_seconds=900), local) + assert _lasts(routed.issue_read_grant(f"s3://{BUCKET}/k").expires_at, 900) + assert _lasts(routed.issue_read_grant(local.put("k", b"v")).expires_at, 1800) + + +def test_a_local_run_leaves_a_custom_stores_own_default_alone(tmp_path): + class _CustomStore(S3ObjectStore): # a store written before lifetimes were a setting + def issue_read_grant(self, object_url, *, expires_in: int = 300): + return super().issue_read_grant(object_url, expires_in=expires_in) + + routed = LocalRunObjectStore(_CustomStore(_StubS3(), BUCKET), GrantingObjectStore(str(tmp_path))) + assert _lasts(routed.issue_read_grant(f"s3://{BUCKET}/k").expires_at, 300) + + +def test_config_toml_sets_the_local_stores_grant_lifetime(monkeypatch, tmp_path): + monkeypatch.delenv("AGENT_ENV_OBJECT_STORE", raising=False) + monkeypatch.delenv("AGENT_ENV_CONFIG", raising=False) + monkeypatch.chdir(tmp_path) + (tmp_path / ".agentenv").mkdir() + (tmp_path / ".agentenv" / "config.toml").write_text( + '[stores.object]\nimpl = "agent_env.store.object_store:LocalFilesystemObjectStore"\n' + 'config = { grant_lifetime_seconds = 3600 }\n' + ) + assert Config().get_object_store().grant_lifetime_seconds == 3600 diff --git a/tst/unit/store/local_grant_server_test.py b/tst/unit/store/local_grant_server_test.py index e16b437e..d88de1a1 100644 --- a/tst/unit/store/local_grant_server_test.py +++ b/tst/unit/store/local_grant_server_test.py @@ -244,6 +244,18 @@ def test_grants_refuse_the_stores_own_files(store): store.issue_read_grant(store.object_url(".agentenv-meta/a")) +def test_local_grants_last_the_stores_lifetime(store, tmp_path): + def lasts(grant, seconds): + return abs(grant.expires_at.timestamp() - time.time() - seconds) < 2 + + assert lasts(store.issue_read_grant(store.put("k", b"v")), 12 * 60 * 60) + short = LocalFilesystemObjectStore( + str(tmp_path / "short"), grant_bind_host="127.0.0.1", grant_advertise_host="localhost", grant_lifetime_seconds=900 + ) + grant = short.issue_read_grant(short.put("k", b"v")) + assert lasts(grant, 900) and httpx.get(grant.url).content == b"v" + + 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) diff --git a/tst/unit/task_step/capture_stubs.py b/tst/unit/task_step/capture_stubs.py index 554c24d2..f5aae855 100644 --- a/tst/unit/task_step/capture_stubs.py +++ b/tst/unit/task_step/capture_stubs.py @@ -80,7 +80,7 @@ def object_url(self, key: str) -> str: def grants_reach(self, sandbox_type: str | None) -> bool: return True - def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in): + def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in=None): self.write_grants.append( { "object_url": object_url, diff --git a/tst/unit/task_step/test_prompt_agent.py b/tst/unit/task_step/test_prompt_agent.py index 2ecf61e3..60263032 100644 --- a/tst/unit/task_step/test_prompt_agent.py +++ b/tst/unit/task_step/test_prompt_agent.py @@ -177,7 +177,7 @@ def object_url(self, key): def grants_reach(self, sandbox_type): return True - def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in): + def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in=None): grants_on_loop.append(_on_loop()) assert object_url == "s3://b/traj/trajectory-CLIENT-ID.json" assert media_type == "application/json" diff --git a/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_trajectory.py b/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_trajectory.py index 03546358..b6605ebc 100644 --- a/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_trajectory.py +++ b/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_trajectory.py @@ -78,7 +78,7 @@ def put(self, key, body, content_type=None): self.puts.append(key) return self.object_url(key) - def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in): + def issue_write_grant(self, object_url, *, media_type, max_bytes, expires_in=None): self.write_grants.append((object_url, media_type, max_bytes)) if self.grant_error is not None: raise self.grant_error diff --git a/tst/util/granting_object_store.py b/tst/util/granting_object_store.py index 1089d00f..52742c31 100644 --- a/tst/util/granting_object_store.py +++ b/tst/util/granting_object_store.py @@ -53,7 +53,8 @@ def get_object_metadata(self, key: str) -> ObjectMetadata | None: return None return replace(metadata, content_type=self._content_types.get(key)) - def issue_read_grant(self, object_url: str, *, expires_in: int = 3600) -> HttpGetGrant: + def issue_read_grant(self, object_url: str, *, expires_in: int | None = None) -> HttpGetGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in return HttpGetGrant( kind="http-get", url=self._grant_url(object_url, "read"), @@ -61,8 +62,9 @@ def issue_read_grant(self, object_url: str, *, expires_in: int = 3600) -> HttpGe ) def issue_write_grant( - self, object_url: str, *, media_type: str, max_bytes: int, expires_in: int = 3600 + self, object_url: str, *, media_type: str, max_bytes: int, expires_in: int | None = None ) -> HttpPutGrant: + expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in return HttpPutGrant( kind="http-put", url=self._grant_url(object_url, "write"),