feat(sync_env_clock): an optional timezone reads virtual_time in the zone and tells every synced server (FD-5003)

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]

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}

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"