feat(prompt_agent): run a multi-turn conversation from one env event to the next on simulated time

src/agent_env/a2a_agent/conversation_store.py

@@ -109,11 +109,20 @@ def find_by_a2a_task_id(a2a_task_id: str) -> Optional[dict]:
     )
 
 
+def _message(role: str, parts: list[dict], now: datetime, env_clock: Optional[dict]) -> dict:
+    """A conversation message; ``env_clock`` maps each env to its virtual time when the message was recorded."""
+    message = {"role": role, "parts": parts, "ts": now.isoformat()}
+    if env_clock is not None:
+        message["env_clock"] = env_clock
+    return message
+
+
 def add_a2a_task(
     conversation_id: str,
     parts: list[dict],
     a2a_task_id: str,
     role: str,
+    env_clock: Optional[dict] = None,
 ) -> Optional[dict]:
     """Append a caller-side message and create a new working a2a_task.
 
@@ -133,7 +142,7 @@ def add_a2a_task(
     doesn't exist.
     """
     now = datetime.now(timezone.utc)
-    new_message = {"role": role, "parts": parts, "ts": now.isoformat()}
+    new_message = _message(role, parts, now, env_clock)
 
     def _mutate(doc: dict) -> UpdateSpec:
         idx = len(doc.get("messages", []))
@@ -161,6 +170,7 @@ def complete_a2a_task(
     conversation_id: str,
     parts: list[dict],
     role: str,
+    env_clock: Optional[dict] = None,
 ) -> Optional[dict]:
     """Append a responder-side message and complete the oldest working a2a_task.
 
@@ -185,7 +195,7 @@ def complete_a2a_task(
     doesn't exist (or the target task lost the state race).
     """
     now = datetime.now(timezone.utc)
-    new_message = {"role": role, "parts": parts, "ts": now.isoformat()}
+    new_message = _message(role, parts, now, env_clock)
     # One backend for the whole retry loop: a reset between the read and the CAS would
     # decide against one database and write to another.
     docs = _doc_store()

src/agent_env/env/gateway/clock.py

@@ -123,6 +123,7 @@ def __init__(self, now_fn=time.monotonic):
         self._rate = 1.0
         self._generation = 0
         self._lived = 0.0  # virtual seconds run before the current arm/clear
+        self._ahead = 0.0  # virtual seconds the current arm was advanced by
 
     @property
     def armed(self) -> bool:
@@ -153,23 +154,41 @@ def set_time(self, virtual_time: str, rate=1.0) -> dict:
             self._bank(t0)
             self._t0 = t0
             self._rate = r
+            self._ahead = 0.0
             self._armed = True
             self._generation += 1
         return self.state()
 
+    def advance(self, virtual_time: str) -> dict:
+        """Jump forward to virtual_time on the current arm: t0, rate and generation stay, so time triggers keep their marks."""
+        target = _parse_rfc3339(virtual_time)
+        with self._lock:
+            if not self._armed or self._t0 is None:
+                raise ClockError("clock is not armed")
+            if target > _MAX_T0 - timedelta(seconds=_OVERFLOW_HORIZON_S * self._rate):
+                raise ClockError(f"virtual_time {virtual_time!r} at rate {self._rate:g} would overflow datetime near the max")
+            current = _advance(self._t0, (self._now() - self._anchor) * self._rate + self._ahead)
+            jump = (target - current).total_seconds()
+            if jump < 0:
+                raise ClockError(f"advance only moves forward: {virtual_time!r} is before the current {_iso(current)}")
+            self._ahead += jump
+            self._lived += jump
+        return self.state()
+
     def clear(self) -> dict:
         """Disarm (idempotent). Afterwards reads 404."""
         with self._lock:
             self._bank()
             self._armed = False
             self._t0 = None
             self._rate = 1.0
+            self._ahead = 0.0
             self._generation += 1
         return {"ok": True, "armed": False}
 
     def elapsed(self) -> float:
         """Virtual seconds run so far: real time at each arm's rate (1 while unarmed), plus the jump of each re-arm
-        forward; one back counts nothing."""
+        or advance forward; one back counts nothing."""
         with self._lock:
             return self._lived + (self._now() - self._anchor) * self._rate
 
@@ -178,7 +197,7 @@ def now(self) -> datetime | None:
             if not self._armed or self._t0 is None:
                 return None
             elapsed = self._now() - self._anchor
-            return _advance(self._t0, elapsed * self._rate)
+            return _advance(self._t0, elapsed * self._rate + self._ahead)
 
     def read(self) -> dict:
         """The frozen consumer contract: {"virtual_time": "<RFC3339>"} — pure read, never advances."""
@@ -193,13 +212,13 @@ def state(self) -> dict:
                 return {"armed": False}
             elapsed = self._now() - self._anchor
             return {"armed": True, "t0": _iso(self._t0), "virtual_seconds_per_real_second": self._rate,
-                    "virtual_time": _iso(_advance(self._t0, elapsed * self._rate))}
+                    "virtual_time": _iso(_advance(self._t0, elapsed * self._rate + self._ahead))}
 
     def _bank(self, t0: datetime | None = None) -> None:
         """Fold the stretch since the last arm/clear into `_lived` at its rate, and start the next one now."""
         now = self._now()
         ran = (now - self._anchor) * self._rate
         self._lived += ran
         if t0 is not None and self._armed:  # a re-arm forward is time the world lived through; one back adds nothing
-            self._lived += max(0.0, (t0 - _advance(self._t0, ran)).total_seconds())
+            self._lived += max(0.0, (t0 - _advance(self._t0, ran + self._ahead)).total_seconds())
         self._anchor = now

src/agent_env/env/gateway/constants.py

@@ -195,10 +195,11 @@ def random_mcp_server_name() -> str:
     },
     {
         "uri": EXT_CLOCK_URI,
-        "description": "Virtual clock (urn:agentenv:clock/v1): read the gateway's virtual time (wall-clock-derived, rate-scaled) via get_time; set-time/clear are the harness control path.",
+        "description": "Virtual clock (urn:agentenv:clock/v1): read the gateway's virtual time (wall-clock-derived, rate-scaled) via get_time; set-time/advance/clear are the harness control path.",
         "params": {"endpoint": "/clock/time", "methods": {
             "get_time": {"method": "GET", "endpoint": "/clock/time"},
             "set_time": {"method": "PUT", "endpoint": "/clock/set-time", "request": {"type": "object", "properties": {"virtual_time": {"type": "string"}, "virtual_seconds_per_real_second": {"type": "number"}}, "required": ["virtual_time"]}},
+            "advance": {"method": "PUT", "endpoint": "/clock/advance", "request": {"type": "object", "properties": {"virtual_time": {"type": "string"}}, "required": ["virtual_time"]}},
             "clear": {"method": "POST", "endpoint": "/clock/clear", "request": {"type": "object", "properties": {}}},
             "state": {"method": "GET", "endpoint": "/clock/state", "request": {"type": "object", "properties": {}}}}},
         "required": False,

src/agent_env/env/gateway/gateway.py

@@ -276,6 +276,20 @@ async def clock_set_time(request: Request) -> Response:
             self._register_get_time_tool()
             return self._json_response(state, 200)
 
+        @self._mcp.custom_route("/clock/advance", methods=["PUT"])
+        async def clock_advance(request: Request) -> Response:
+            try:
+                body = await request.json()
+            except Exception:
+                return self._json_response({"ok": False, "error": "Invalid JSON body"}, 400)
+            if not isinstance(body, dict):
+                return self._json_response({"ok": False, "error": "body must be a JSON object"}, 400)
+            try:
+                state = self._clock.advance(body.get("virtual_time"))
+            except ClockError as e:
+                return self._json_response({"ok": False, "error": str(e)}, 400)
+            return self._json_response(state, 200)
+
         @self._mcp.custom_route("/clock/clear", methods=["POST"])
         async def clock_clear(request: Request) -> Response:
             self._unregister_get_time_tool()