diff --git a/pyproject.toml b/pyproject.toml index 397447cb6..1be42de95 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ dependencies = [ "modal>=1.5.2.dev22", # dev22 adds _experimental_create(tags=); dev2 added VM-mode V2 (vm_runtime) "e2b~=2.46.4", # AsyncSandbox/AsyncTemplate APIs validated against the 2.46 line "fastapi>=0.110", # the explorer, agent-env up + "tzdata>=2024.1", # sync_env_clock reads a task's timezone on a worker image with no system zone files ] [project.optional-dependencies] diff --git a/src/agent_env/task_step/task_steps/sync_env_clock.py b/src/agent_env/task_step/task_steps/sync_env_clock.py index ede3a409b..76ff196e8 100644 --- a/src/agent_env/task_step/task_steps/sync_env_clock.py +++ b/src/agent_env/task_step/task_steps/sync_env_clock.py @@ -9,9 +9,12 @@ from __future__ import annotations import logging +from datetime import datetime, timezone as _utc_tz from typing import ClassVar, Optional +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import httpx +import tzdata # noqa: F401 from agentenv_protocol import client as protocol_v1 from agent_env.env.gateway.constants import EXT_CLOCK_URI @@ -22,6 +25,32 @@ logger = logging.getLogger(__name__) +def resolve_virtual_time(virtual_time: str, timezone: Optional[str]) -> str: + """The instant to arm the gateway at. Without a zone, `virtual_time` unchanged; with one, a naive wall + time is read in that zone and an explicit offset must agree with it. Returned as a UTC RFC3339 string.""" + if timezone is None: + return virtual_time + try: + zone = ZoneInfo(timezone.strip()) if isinstance(timezone, str) and timezone.strip() else None + except (ZoneInfoNotFoundError, ValueError, OSError): + zone = None + if zone is None: + raise ValueError(f"timezone must be a valid IANA name such as 'America/Chicago', got {timezone!r}") + try: + parsed = datetime.fromisoformat(virtual_time) + except ValueError as e: + raise ValueError(f"virtual_time {virtual_time!r} is not an ISO 8601 / RFC3339 timestamp") from e + if parsed.tzinfo is None: + local = parsed.replace(tzinfo=zone) + if local.astimezone(_utc_tz.utc).astimezone(zone).replace(tzinfo=None) != parsed: + raise ValueError(f"virtual_time {virtual_time!r} does not exist in {zone.key} (clock skips it)") + return local.astimezone(_utc_tz.utc).isoformat().replace("+00:00", "Z") + if parsed.astimezone(zone).utcoffset() != parsed.utcoffset(): + raise ValueError(f"virtual_time {virtual_time!r} carries offset {parsed.utcoffset()} but {zone.key} is at " + f"{parsed.astimezone(zone).utcoffset()} at that instant; drop the offset or fix the timezone") + return parsed.astimezone(_utc_tz.utc).isoformat().replace("+00:00", "Z") + + class SyncEnvClockTaskStep(TaskStep): type: ClassVar[str] = "sync_env_clock" entity_refs = (EntityRef.env("env_id"),) @@ -37,9 +66,13 @@ def __init__( depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True, timeout_seconds: int = 30, + timezone: Optional[str] = None, ): """virtual_seconds_per_real_second = clock speed (1.0 real time, 0 frozen, 86400 cap = - 1 real s -> 1 virtual day). For ~V virtual seconds over an ~R-real-second task, pass V/R.""" + 1 real s -> 1 virtual day). For ~V virtual seconds over an ~R-real-second task, pass V/R. + + timezone = optional IANA name. Unset keeps today's behaviour. Set, a wall-clock virtual_time is local time in + that zone, and every synced server is told the zone (a server that cannot take it fails the step).""" super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) if not isinstance(virtual_time, str): raise ValueError("virtual_time must be an RFC3339 string") @@ -48,6 +81,8 @@ def __init__( self.virtual_seconds_per_real_second = virtual_seconds_per_real_second self.tolerate_missing_sync_time = tolerate_missing_sync_time self.timeout_seconds = timeout_seconds + self.timezone = timezone.strip() if isinstance(timezone, str) else timezone + self.armed_virtual_time = resolve_virtual_time(virtual_time, self.timezone) def to_dict(self) -> dict: base = super().to_dict() @@ -56,6 +91,8 @@ def to_dict(self) -> dict: base["virtual_seconds_per_real_second"] = self.virtual_seconds_per_real_second base["tolerate_missing_sync_time"] = self.tolerate_missing_sync_time base["timeout_seconds"] = self.timeout_seconds + if self.timezone is not None: + base["timezone"] = self.timezone return base @classmethod @@ -67,6 +104,7 @@ def from_dict(cls, data: dict) -> SyncEnvClockTaskStep: virtual_seconds_per_real_second=data.get("virtual_seconds_per_real_second", 1.0), tolerate_missing_sync_time=data.get("tolerate_missing_sync_time", True), timeout_seconds=data.get("timeout_seconds", 30), + timezone=data.get("timezone"), ) def _service_names(self, deployed) -> list[str]: @@ -86,7 +124,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") try: await deployed.invoke(EXT_CLOCK_URI, "set_time", - {"virtual_time": self.virtual_time, "virtual_seconds_per_real_second": self.virtual_seconds_per_real_second}, + {"virtual_time": self.armed_virtual_time, "virtual_seconds_per_real_second": self.virtual_seconds_per_real_second}, timeout=self.timeout_seconds) except httpx.HTTPStatusError as e: raise RuntimeError(f"clock arm failed (HTTP {e.response.status_code}): {e.response.text}") from e @@ -98,11 +136,14 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: outcome = await self._sync_one(deployed, gateway_url_of(deployed), service, env_get_time_url) (synced if outcome.get("synced") else skipped).append({"service": service, "environment": service, **outcome}) - context.metadata.setdefault("clock_configurations", []).append({ + configuration = { "step_id": self.id, "env_id": self.env_id, "virtual_time": self.virtual_time, "virtual_seconds_per_real_second": self.virtual_seconds_per_real_second, "env_get_time_url": env_get_time_url, "synced": synced, "skipped": skipped, - }) + } + if self.timezone is not None: + configuration.update(timezone=self.timezone, armed_at=self.armed_virtual_time) + context.metadata.setdefault("clock_configurations", []).append(configuration) logger.info(f"sync_env_clock env={self.env_id}: armed rate={self.virtual_seconds_per_real_second}; " f"synced={[s['service'] for s in synced]} skipped={[s['service'] for s in skipped]}") if not synced: @@ -126,7 +167,14 @@ async def _sync_one(self, deployed, gateway_url: str, service: str, env_get_time raise RuntimeError(f"{service} does not advertise {EXT_CLOCK_URI}") if not env_get_time_url: raise RuntimeError(f"gateway did not expose env_get_time_url; cannot sync {service}") + params = {"env_get_time_url": env_get_time_url} + if self.timezone is not None: + method = protocol_v1.find_extension_method(card, EXT_CLOCK_URI, "sync_time") or {} + if "timezone" not in ((method.get("request") or {}).get("properties") or {}): + raise RuntimeError(f"{service} advertises {EXT_CLOCK_URI} but its sync_time takes no timezone; " + f"rebuild its image or drop timezone={self.timezone!r} from the task") + params["timezone"] = self.timezone result = await protocol_v1.invoke_extension(base_url, card, EXT_CLOCK_URI, - params={"env_get_time_url": env_get_time_url}, + params=params, timeout=self.timeout_seconds, headers=deployed.control_headers) return {"synced": True, "result": result} diff --git a/tst/unit/task_step/sync_env_clock_test.py b/tst/unit/task_step/sync_env_clock_test.py index 471ca3a3a..5ab7bee8e 100644 --- a/tst/unit/task_step/sync_env_clock_test.py +++ b/tst/unit/task_step/sync_env_clock_test.py @@ -7,7 +7,7 @@ from agent_env.env.gateway.constants import EXT_CLOCK_URI from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps import sync_env_clock as mod -from agent_env.task_step.task_steps.sync_env_clock import SyncEnvClockTaskStep +from agent_env.task_step.task_steps.sync_env_clock import SyncEnvClockTaskStep, resolve_virtual_time T0 = "2026-01-01T00:00:00Z" CLOCK_URL = "http://gateway:9000/clock/time" @@ -125,6 +125,10 @@ async def test_execute_keeps_the_arm_failure_message(monkeypatch): _CLOCK_EXT = {"uri": EXT_CLOCK_URI, "params": {"endpoint": "/agentenv/ext/sync_time", "methods": {"sync_time": {"method": "POST"}}}} +_ZONE_EXT = {"uri": EXT_CLOCK_URI, "params": {"endpoint": "/agentenv/ext/sync_time", "methods": {"sync_time": {"method": "POST", "request": { + "type": "object", "properties": {"env_get_time_url": {"type": "string"}, "timezone": {"type": "string"}}, "required": ["env_get_time_url"]}}}}} +_OLD_EXT = {"uri": EXT_CLOCK_URI, "params": {"endpoint": "/agentenv/ext/sync_time", "methods": {"sync_time": {"method": "POST", "request": { + "type": "object", "properties": {"env_get_time_url": {"type": "string"}}, "required": ["env_get_time_url"]}}}}} # The card's address differs from the gateway URL, so a call built from the stored card is told apart from one built from /svc/... itself. _CARD = "https://sandbox.example/sb-1" @@ -171,3 +175,95 @@ def handle(request: httpx.Request) -> httpx.Response: monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) return sent + + +def _zoned(virtual_time="2026-01-01T12:00:00", timezone="America/New_York", **kw): + return SyncEnvClockTaskStep(id="sync", version=1, env_id="e1", virtual_time=virtual_time, timezone=timezone, depends_on=[], **kw) + + +def test_timezone_round_trips_and_is_absent_from_an_unzoned_step(): + d = _zoned().to_dict() + assert d["timezone"] == "America/New_York" + assert SyncEnvClockTaskStep.from_dict(d).timezone == "America/New_York" + assert "timezone" not in _step().to_dict() + assert SyncEnvClockTaskStep.from_dict(_step().to_dict()).timezone is None + + +@pytest.mark.parametrize("name", ["", " ", "Mars/Olympus", 7, "Europe", "x" * 5000, "../x"]) +def test_an_invalid_zone_is_refused_when_the_step_is_built(name): + with pytest.raises(ValueError, match="IANA"): + _zoned(timezone=name) + + +@pytest.mark.parametrize("virtual_time, zone, armed", [ + ("2026-01-01T12:00:00", "America/New_York", "2026-01-01T17:00:00Z"), + ("2026-07-01T12:00:00", "America/New_York", "2026-07-01T16:00:00Z"), + ("2026-01-01T12:00:00", "UTC", "2026-01-01T12:00:00Z"), + ("2026-01-01T12:00:00-05:00", "America/New_York", "2026-01-01T17:00:00Z"), + ("2026-01-01T17:00:00Z", "UTC", "2026-01-01T17:00:00Z"), + ("2026-11-01T01:30:00", "America/New_York", "2026-11-01T05:30:00Z"), +]) +def test_a_wall_clock_time_is_read_in_the_zone(virtual_time, zone, armed): + assert _zoned(virtual_time, zone).armed_virtual_time == armed + + +@pytest.mark.parametrize("virtual_time, zone, match", [ + ("2026-01-01T12:00:00Z", "America/New_York", "carries offset"), + ("2026-07-01T12:00:00-05:00", "America/New_York", "carries offset"), + ("2026-03-08T02:30:00", "America/New_York", "does not exist"), + ("not a time", "America/New_York", "not an ISO"), +]) +def test_a_time_that_disagrees_with_the_zone_is_refused(virtual_time, zone, match): + with pytest.raises(ValueError, match=match): + _zoned(virtual_time, zone) + + +def test_without_a_zone_virtual_time_is_passed_through_untouched(): + assert _step().armed_virtual_time == T0 + assert resolve_virtual_time("2026-01-01T12:00:00", None) == "2026-01-01T12:00:00" + + +@pytest.mark.asyncio +async def test_execute_arms_the_gateway_at_the_local_instant_and_tells_every_server_the_zone(monkeypatch): + sent = _mock_http(monkeypatch, state={"env_get_time_url": CLOCK_URL}) + monkeypatch.setattr(SyncEnvClockTaskStep, "_service_names", lambda self, deployed: ["email"]) + context = TaskStepContext(deployed_envs=[_carded(await _composed_card({"mcp-email": [_ZONE_EXT]}))]) + + await _zoned().execute(context) + + assert json.loads(sent[0].content)["virtual_time"] == "2026-01-01T17:00:00Z" + assert json.loads(sent[2].content) == {"env_get_time_url": CLOCK_URL, "timezone": "America/New_York"} + cfg = context.metadata["clock_configurations"][0] + assert (cfg["timezone"], cfg["virtual_time"], cfg["armed_at"]) == ("America/New_York", "2026-01-01T12:00:00", "2026-01-01T17:00:00Z") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tolerate", [True, False]) +async def test_a_server_that_cannot_take_the_zone_fails_the_step_before_any_sync_is_sent(monkeypatch, tolerate): + sent = _mock_http(monkeypatch, state={"env_get_time_url": CLOCK_URL}) + monkeypatch.setattr(SyncEnvClockTaskStep, "_service_names", lambda self, deployed: ["email"]) + context = TaskStepContext(deployed_envs=[_carded(await _composed_card({"mcp-email": [_OLD_EXT]}))]) + + with pytest.raises(RuntimeError, match="takes no timezone"): + await _zoned(tolerate_missing_sync_time=tolerate).execute(context) + assert not any(str(r.url).endswith("/sync_time") for r in sent) + + +@pytest.mark.asyncio +async def test_an_unzoned_step_sends_the_same_body_to_an_old_and_a_new_server(monkeypatch): + sent = _mock_http(monkeypatch) + for ext in (_OLD_EXT, _ZONE_EXT, _CLOCK_EXT): + await _step()._sync_one(_carded(await _composed_card({"mcp-email": [ext]})), "http://gw", "email", CLOCK_URL) + assert [json.loads(r.content) for r in sent] == [{"env_get_time_url": CLOCK_URL}] * 3 + + +@pytest.mark.asyncio +async def test_a_server_that_does_not_advertise_clock_is_still_skipped_when_a_zone_is_set(monkeypatch): + _mock_http(monkeypatch) + out = await _zoned()._sync_one(_carded(await _composed_card({"mcp-email": []})), "http://gw", "email", CLOCK_URL) + assert out == {"synced": False, "reason": "not_advertised"} + + +def test_a_padded_zone_is_stored_and_sent_stripped(): + s = _zoned(timezone=" America/New_York ") + assert s.timezone == "America/New_York" and s.to_dict()["timezone"] == "America/New_York" diff --git a/uv.lock b/uv.lock index 8e3df6800..18b83db83 100644 --- a/uv.lock +++ b/uv.lock @@ -68,6 +68,7 @@ dependencies = [ { name = "python-multipart" }, { name = "pyyaml" }, { name = "starlette" }, + { name = "tzdata" }, { name = "uvicorn" }, ] @@ -251,6 +252,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'skyrl'", specifier = ">=5" }, { name = "transformers", marker = "extra == 'trl'", specifier = ">=5.2" }, { name = "transformers", marker = "extra == 'verl'", specifier = ">=5" }, + { name = "tzdata", specifier = ">=2024.1" }, { name = "urllib3", marker = "extra == 'boat'", specifier = ">=2.1.0,<3" }, { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" },