feat(gateway)!: clock writes take the control token, a watch its computer's read token, and a waiting write names a computer that stopped watching (FD-5472)

src/agent_env/env/gateway/clock_followers.py

@@ -32,6 +32,9 @@ class ClockFollowers:
     run on the loop that serves it, never one asyncio.run apiece."""
 
     WATCH_S = CLOCK_WATCH_HOLD_S
+    #: How long since a follower last watched before it has gone quiet: longer than a watch is held and a follower takes
+    #: to reconnect, so its connection is gone. A write doesn't wait for a quiet follower; it names it at once.
+    QUIET_S = CLOCK_WATCH_HOLD_S + 5
 
     def __init__(self, clock: Callable[[], Clock], tokens: Mapping[str, str]):
         self._clock = clock
@@ -51,6 +54,23 @@ def follower_for(self, name: str | None, token: str) -> str | None:
         expected = self._tokens.get(name or "")
         return name if expected and token and hmac.compare_digest(token.encode(), expected.encode()) else None
 
+    def holder_of(self, token: str) -> str | None:
+        """The declared follower whose read token ``token`` is, else None."""
+        found = None
+        for name, expected in self._tokens.items():
+            if token and hmac.compare_digest(token.encode(), expected.encode()):
+                found = name
+        return found
+
+    def _quiet_at(self, name: str) -> float | None:
+        """When ``name`` goes quiet (monotonic), if it has watched."""
+        seen = self._seen.get(name)
+        return None if seen is None else seen + self.QUIET_S
+
+    def _quiet(self, name: str) -> bool:
+        at = self._quiet_at(name)
+        return at is not None and time.monotonic() >= at
+
     async def watch(self, after: str, follower: str | None) -> None:
         """Take ``after`` as ``follower``'s acknowledgement (one ``follower_for`` named) when it's a version this clock
         has had, later than the one it last showed; then return at once when the clock's version isn't ``after``, else
@@ -71,18 +91,23 @@ async def watch(self, after: str, follower: str | None) -> None:
                 pass
 
     async def stale_after(self, version: str, wait_s: float) -> list[str]:
-        """The followers that haven't shown ``version`` (or a later one) once each has or ``wait_s`` has passed."""
+        """The followers that haven't shown ``version`` (or a later one) once each has, or has gone quiet (``QUIET_S``),
+        or ``wait_s`` has passed: woken by each acknowledgement and when the next follower waited for goes quiet. One
+        that hasn't watched yet is waited for: it may be starting."""
         target = self._clock().count_of(version)
+        deadline = time.monotonic() + wait_s
 
         def stale() -> list[str]:
             return sorted(name for name, shown in self._shown.items()
                           if shown is None or target is None or shown < target)
 
         async with self._acks:
-            try:
-                await asyncio.wait_for(self._acks.wait_for(lambda: not stale()), wait_s)
-            except asyncio.TimeoutError:
-                pass
+            while (waited_for := [name for name in stale() if not self._quiet(name)]) and time.monotonic() < deadline:
+                wake = min([deadline, *filter(None, map(self._quiet_at, waited_for))])
+                try:
+                    await asyncio.wait_for(self._acks.wait(), max(0.0, wake - time.monotonic()))
+                except asyncio.TimeoutError:
+                    pass
             return stale()
 
     def state(self) -> dict:

src/agent_env/env/gateway/constants.py

@@ -168,14 +168,18 @@ def sources(names: list[str]) -> str:
 WILDCARD = "*"
 TOOL_DISABLE_ACTION = "disable"
 TOOL_ENABLE_ACTION = "enable"
-# The harness's routes that take the deploy's control token: the ones that change tool access, triggers and
-# notifications. Every other route stays open for now, the clock's writes included: callers outside agent-env set the
-# clock over plain HTTP and don't send the token yet.
+# The harness's routes that take the deploy's control token: the ones that change tool access, triggers, notifications
+# and the clock. Its reads stay open (the clock's state and time show only the universe's time, its rate and which
+# version each computer shows), but for the watch, CLOCK_READ_ROUTES.
 CONTROL_ROUTES = frozenset({
     f"/tools/{TOOL_DISABLE_ACTION}", f"/tools/{TOOL_ENABLE_ACTION}",
     "/triggers/register", "/triggers/remove", "/triggers/clear",
     "/notifications/config", "/notifications/take", "/notifications/clear",
+    "/clock/set-time", "/clock/set-rate", "/clock/advance", "/clock/clear",
 })
+# The routes that take the deploy's control token or a declared computer's read token: the watch, which is a
+# computer's acknowledgement of the version it names, so a watch carrying a computer's token names that computer.
+CLOCK_READ_ROUTES = frozenset({"/clock/watch"})
 
 
 class GatewayMode(str, Enum):

src/agent_env/env/gateway/gateway.py

@@ -48,6 +48,7 @@
     CHANGES_PAGE_DEFAULT,
     CHANGES_PAGE_MAX,
     CHILD_DISCOVERY_TIMEOUT_S,
+    CLOCK_READ_ROUTES,
     CONTROL_ROUTES,
     DATA_PLANE_LOAD_TIMEOUT_MAX_S,
     DEFAULT_ROLE,
@@ -472,6 +473,10 @@ def _asgi_app(self):
         async def app(scope, receive, send):
             if self._control_bearer and scope["type"] in ("http", "websocket"):
                 headers = scope.get("headers", ())
+                if scope["type"] == "http" and scope["path"] in CLOCK_READ_ROUTES and (
+                        refusal := self._clock_read_refusal(scope)) is not None:
+                    await refusal(scope, receive, send)
+                    return
                 if scope["type"] == "http" and scope["path"] in CONTROL_ROUTES and not hmac.compare_digest(
                         dict(headers).get(b"authorization", b""), self._control_bearer):
                     response = self._json_response({"ok": False, "error": "this route takes the env's control token"}, 401)
@@ -496,6 +501,26 @@ async def app(scope, receive, send):
 
         return app
 
+    def _clock_read_refusal(self, scope) -> Response | None:
+        """Why a read of the clock is refused, or None when it may go on. It takes the gateway's control token or a
+        declared computer's read token, and a watch carrying a computer's token names that computer (a watch is its
+        acknowledgement of the version it names); 401 without a token, 403 for one that's neither or names another."""
+        authorization = dict(scope.get("headers", ())).get(b"authorization", b"")
+        if hmac.compare_digest(authorization, self._control_bearer):
+            return None
+        if not authorization.startswith(b"Bearer "):
+            return self._json_response({"ok": False, "error": (
+                "the clock's reads take the env's control token or a computer's read token")}, 401)
+        holder = self._clock_followers.holder_of(authorization[len(b"Bearer "):].decode(errors="replace").strip())
+        if holder is None:
+            return self._json_response({"ok": False, "error": "the token is no computer's read token"}, 403)
+        if scope["path"] == "/clock/watch":
+            named = dict(parse_qsl(scope.get("query_string", b"").decode(errors="replace"))).get("follower", "")
+            if named != holder:
+                return self._json_response({"ok": False, "error": (
+                    f"a watch with {holder}'s read token names {holder} as its follower, not {named or 'none'}")}, 403)
+        return None
+
     @asynccontextmanager
     async def _session_lifespan(self, mcp: FastMCP):
         """Once per inbound MCP session: the session's key joins the caller, and a connect gets to heal a failed
@@ -612,7 +637,8 @@ async def _await_barriers(self, pending: list[asyncio.Task], timeout_s: float |
 
     async def _clock_write(self, request: Request, change, *, body_optional: bool = False) -> Response:
         """A clock write route: its JSON object body (none at all when ``body_optional``) applied by ``change`` through
-        ``_change_clock``; 400 for a body or value the clock refuses, else 200, or 504 naming the followers behind."""
+        ``_change_clock``; 400 for a body or value the clock refuses, else 200, or 504 naming the followers behind, with
+        what each follower last showed and when it last watched."""
         try:
             body = await request.json() if not body_optional or await request.body() else {}
         except Exception:
@@ -625,9 +651,10 @@ async def _clock_write(self, request: Request, change, *, body_optional: bool =
         except ClockError as e:
             return self._json_response({"ok": False, "error": str(e)}, 400)
         if stale:
-            return self._json_response({**state, "ok": False, "stale": stale, "error": (
-                f"the clock is at version {state['version']}, but {', '.join(stale)} didn't show it within "
-                f"{wait:g} s: their clocks are behind")}, 504)
+            return self._json_response({**state, "ok": False, "stale": stale, "followers": self._clock_followers.state(),
+                                        "error": (f"the clock is at version {state['version']}, but {', '.join(stale)} "
+                                                  f"didn't show it (waited up to {wait:g} s; one that has stopped "
+                                                  "watching isn't waited for): their clocks are behind")}, 504)
         return self._json_response(state, 200)
 
     async def _change_clock(self, change, wait_s: float = 0.0) -> tuple[dict, list[str]]:

tst/integration/computer/multi_env_computers_test.py

@@ -5,12 +5,15 @@
 and a role rule narrows them as it would a server's; /svc/<computer>/ reaches the computer. A later process reattaches the
 computers from the multi env's record, and the run's teardown leaves nothing up. Two computers keep serving through a
 computer's restart and the gateway's, a failed redeploy leaves the earlier deployment serving, and a computer that fails
-mid-deploy, or an entry the multi env refuses, leaves nothing up. Needs Docker and the local registry on localhost:5000."""
+mid-deploy, or an entry the multi env refuses, leaves nothing up. The gateway's control token, which moves its clock,
+appears nowhere on a computer: in no file, and in no process's environment or command line. Needs Docker and the local
+registry on localhost:5000."""
 
 import asyncio
 import json
 import re
 import shutil
+import subprocess
 import tempfile
 import uuid
 from pathlib import Path
@@ -166,6 +169,41 @@ def _containers(sandbox_id: str) -> list[str]:
     return _docker("ps", "-aq", "--filter", f"name=agent-{sandbox_id}").stdout.split()
 
 
+# Run in a computer: the files and the processes' environments and command lines holding the token read from stdin, as
+# far as the user it runs as can read them (so the token is never in its own environment or command line).
+TRACES = r"""import os, sys
+LARGEST_FILE_BYTES = 50_000_000  # past this a file is media or a database, not a config or a secret
+token = sys.stdin.read().strip().encode()
+found = []
+for root in ("/etc", "/opt", "/home", "/root", "/tmp", "/run", "/var", "/srv", "/usr/local"):
+    for folder, _, names in os.walk(root):
+        for name in names:
+            path = os.path.join(folder, name)
+            try:
+                if os.path.isfile(path) and os.path.getsize(path) < LARGEST_FILE_BYTES and token in open(path, "rb").read():
+                    found.append(path)
+            except OSError:
+                pass
+for pid in filter(str.isdigit, os.listdir("/proc")):
+    for part in ("environ", "cmdline"):
+        try:
+            if token in open(f"/proc/{pid}/{part}", "rb").read():
+                found.append(f"/proc/{pid}/{part}")
+        except OSError:
+            pass
+print("\n".join(found))
+"""
+
+
+def _traces(sandbox_id: str, token: str) -> list[str]:
+    """Where in the computer's container ``token`` is, read as root and as the session user (uid 1000), whose processes'
+    environments root can't read there."""
+    (container,) = _containers(sandbox_id)
+    return [line for user in ("0", "1000") for line in subprocess.run(
+        ["docker", "exec", "-i", "-u", user, container, "python3", "-c", TRACES], input=token, capture_output=True,
+        text=True, check=True).stdout.split()]
+
+
 @pytest.fixture
 def made(monkeypatch) -> list[str]:
     return local_sandboxes.recorded_sandbox_ids(monkeypatch)
@@ -226,11 +264,13 @@ async def test_a_computer_in_a_multi_env_is_served_through_its_gateway_under_its
             narrowed, narrowed_shot = await _tools(deployed)
             reattached = await Env.from_instance_id(deployed.instance_id)
             up = _containers(computer["sandbox_id"])
+            traces = _traces(computer["sandbox_id"], deployed.control_token)
 
         finally:
             report = await teardown_run(context)
 
     assert not report.failed
+    assert deployed.control_token and traces == [], traces
     assert deployed.get_child_env_card("desktop") is not None and card["name"] == "desktop"
     assert "list_items" in listed and "desktop_computer_screenshot" in listed and shot
     assert {name for name in listed if "computer_" in name} <= {f"desktop_{tool}" for tool in COMPUTER_TOOLS}

tst/unit/env/gateway/clock_watch_test.py

@@ -1,4 +1,5 @@
-"""The /clock/watch long poll computers follow the gateway's clock through, and clock writes that wait for them."""
+"""The /clock/watch long poll computers follow the gateway's clock through, each watch under its computer's read token
+on a gateway with a control token, and clock writes that wait for them, but not for one that has stopped watching."""
 from __future__ import annotations
 
 import asyncio
@@ -18,6 +19,7 @@
 
 T0 = "2026-01-01T00:00:00Z"
 TOKENS = {"desktop": "token-desktop", "laptop": "token-laptop"}
+CONTROL = "token-control"
 HELD_S = 5.0  # long enough that an answer "at once" can't be a held watch timing out, on a loaded runner too
 SHORT_HOLD_S = 0.3
 
@@ -32,9 +34,9 @@ def gateway(tmp_path, monkeypatch):
     monkeypatch.setattr(ClockFollowers, "WATCH_S", HELD_S)
     made = []
 
-    def make(followers=TOKENS) -> Gateway:
+    def make(followers=TOKENS, control_token=None) -> Gateway:
         made.append(Gateway(host="127.0.0.1", port=0, server_name="unit", internal_mcp_servers=[],
-                            clock_followers=followers))
+                            clock_followers=followers, control_token=control_token))
         return made[-1]
 
     yield make
@@ -165,7 +167,8 @@ async def test_a_write_answers_504_naming_the_computers_still_behind_and_the_cha
     body = answer.json()
     assert answer.status_code == 504 and waited >= 0.4
     assert (body["ok"], body["stale"], body["version"]) == (False, ["laptop"], gw._clock.version)
-    assert "laptop didn't show it within 0.4 s" in body["error"]
+    assert "laptop didn't show it (waited up to 0.4 s;" in body["error"]
+    assert body["followers"]["laptop"] == {"version": None, "seen_s_ago": None}
     assert read == {"virtual_time": T0}
 
 
@@ -363,3 +366,77 @@ async def test_a_change_on_the_arm_wakes_the_triggers_and_a_re_arm_leaves_them_t
         advanced = (await client.put("/clock/advance", json={"virtual_time": "2026-02-01T00:00:00Z"})).json()
         cleared = (await client.post("/clock/clear")).json()
     assert ["due" in answer for answer in (armed, rated, advanced, cleared)] == [False, True, True, False]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("token, follower, status", [
+    ("token-laptop", "laptop", 200), ("token-laptop", "desktop", 403), ("token-laptop", None, 403),
+    ("wrong", "laptop", 403), (None, "laptop", 401), (CONTROL, None, 200), (CONTROL, "laptop", 200)],
+    ids=["its-own", "naming-another", "naming-none", "unknown", "none", "control", "control-naming-it"])
+async def test_a_guarded_gateways_watch_takes_its_computers_read_token_under_its_name(gateway, monkeypatch, token,
+                                                                                         follower, status):
+    """With a control token, a watch takes it or a computer's read token, and a computer's names that computer; only
+    the computer's own acknowledges, and the control token reads without acknowledging."""
+    monkeypatch.setattr(ClockFollowers, "WATCH_S", SHORT_HOLD_S)
+    gw = gateway(control_token=CONTROL)
+    async with _client(gw) as client:
+        armed = await client.put("/clock/set-time", json={"virtual_time": T0}, headers={"Authorization": f"Bearer {CONTROL}"})
+        answered = await _watch(client, gw._clock.version, follower, token)
+    assert armed.status_code == 200 and answered.status_code == status, answered.text
+    own = (token, follower) == ("token-laptop", "laptop")
+    assert gw._clock_followers.state()["laptop"]["version"] == (gw._clock.version if own else None)
+    if (token, follower) == ("token-laptop", "desktop"):
+        assert answered.json()["error"] == "a watch with laptop's read token names laptop as its follower, not desktop"
+
+
+@pytest.mark.asyncio
+async def test_a_write_doesnt_wait_for_a_computer_that_has_stopped_watching(gateway, monkeypatch):
+    """A computer whose last watch is older than a hold and a reconnect (its connection gone) is named at once, with
+    when it last watched; the others are waited for as before."""
+    monkeypatch.setattr(ClockFollowers, "WATCH_S", SHORT_HOLD_S)
+    monkeypatch.setattr(ClockFollowers, "QUIET_S", SHORT_HOLD_S + 0.5)
+    gw = gateway()
+    async with _client(gw) as client:
+        await client.put("/clock/set-time", json={"virtual_time": T0})
+        follower = asyncio.create_task(_follow(client, "desktop", TOKENS["desktop"]))
+        try:
+            await _watch(client, gw._clock.version, "laptop", TOKENS["laptop"])  # shows it, then stops watching
+            await _until_shown(gw, "desktop", "laptop")
+            await asyncio.sleep(ClockFollowers.QUIET_S + 0.2)
+            started = time.monotonic()
+            answer = await client.put("/clock/advance", json={"virtual_time": "2026-01-02T00:00:00Z", "wait_followers_s": 10})
+            waited = time.monotonic() - started
+        finally:
+            follower.cancel()
+    body = answer.json()
+    assert answer.status_code == 504 and body["stale"] == ["laptop"] and waited < 3, (waited, body)
+    assert body["followers"]["laptop"]["seen_s_ago"] > ClockFollowers.QUIET_S
+    assert body["followers"]["desktop"]["version"] == gw._clock.version
+
+
+@pytest.mark.asyncio
+async def test_a_write_waits_for_a_computer_that_hasnt_watched_yet(gateway):
+    """One that hasn't watched may be starting: the write waits for it, all its wait_followers_s."""
+    gw = gateway({"laptop": TOKENS["laptop"]})
+    async with _client(gw) as client:
+        started = time.monotonic()
+        answer = await client.put("/clock/set-time", json={"virtual_time": T0, "wait_followers_s": 0.5})
+        waited = time.monotonic() - started
+    assert answer.status_code == 504 and answer.json()["stale"] == ["laptop"] and waited >= 0.5
+
+
+@pytest.mark.asyncio
+async def test_a_write_stops_waiting_for_a_computer_when_it_goes_quiet_during_the_wait(gateway, monkeypatch):
+    """One that watched shortly before the write is waited for until its quiet deadline, not for the whole wait."""
+    monkeypatch.setattr(ClockFollowers, "WATCH_S", SHORT_HOLD_S)
+    monkeypatch.setattr(ClockFollowers, "QUIET_S", 1.0)
+    gw = gateway({"laptop": TOKENS["laptop"]})
+    async with _client(gw) as client:
+        await client.put("/clock/set-time", json={"virtual_time": T0})
+        await _watch(client, "", "laptop", TOKENS["laptop"])  # answered at once: laptop watched, then stops
+        await asyncio.sleep(0.3)
+        started = time.monotonic()
+        answer = await client.put("/clock/advance", json={"virtual_time": "2026-01-02T00:00:00Z", "wait_followers_s": 10})
+        waited = time.monotonic() - started
+    assert answer.status_code == 504 and answer.json()["stale"] == ["laptop"], answer.text
+    assert 0.5 <= waited < 3, waited