From e89a5399e6fb891758ccdd43282d3f221274f407 Mon Sep 17 00:00:00 2001 From: Sarah Chitty Date: Mon, 26 Oct 2026 23:54:13 -0700 Subject: [PATCH 1/4] =?UTF-8?q?feat(gateway):=20wakeups=20=E2=80=94=20chec?= =?UTF-8?q?ks=20an=20agent=20schedules=20for=20itself,=20and=20notify=5Fag?= =?UTF-8?q?ent?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit What wakes an agent between its turns, on the env's virtual clock (urn:agentenv:wakeups/v1): - schedule_check / list_scheduled_checks / cancel_scheduled_check: gateway built-in tools, next to get_time, that exist only once a task turns them on (configure_wakeups step, or POST /wakeups/config). A check takes `at` (RFC 3339 or a duration from now) and/or `every` with `until`; max_active and min_interval cap them. - notify_agent: a trigger action that queues a notification ({channel, text}, text templated from the provoking call) at the current virtual time. Only task designers write it; the agent has no general alert tool. - GET /wakeups/state says when the next check is due and what is pending; POST /wakeups/take hands out each due check and notification once (a repeating check passed several times in one jump comes out once, with a count). Nothing here wakes anyone: a post-turn hook reads state, moves the clock and takes what is due. Opt-in throughout: without configure_wakeups no tool appears, and a gateway without notify_agent is refused by register_env_triggers before any request. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/constants.py | 17 +- src/agent_env/env/gateway/gateway.py | 86 +++++- src/agent_env/env/gateway/triggers.py | 22 +- src/agent_env/env/gateway/wakeups.py | 286 +++++++++++++++++ src/agent_env/task/episode/rules.py | 3 +- src/agent_env/task_step/registry.py | 2 + .../task_step/task_steps/configure_wakeups.py | 90 ++++++ .../task_steps/register_env_triggers.py | 5 +- tst/unit/env/gateway/gateway_card_test.py | 4 +- tst/unit/env/gateway/wakeups_test.py | 289 ++++++++++++++++++ tst/unit/task_step/test_configure_wakeups.py | 79 +++++ .../task_step/test_register_env_triggers.py | 15 + 12 files changed, 890 insertions(+), 8 deletions(-) create mode 100644 src/agent_env/env/gateway/wakeups.py create mode 100644 src/agent_env/task_step/task_steps/configure_wakeups.py create mode 100644 tst/unit/env/gateway/wakeups_test.py create mode 100644 tst/unit/task_step/test_configure_wakeups.py diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index 4ab0bf77a..75348e187 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -78,6 +78,7 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int: EXT_STEP_URI = "urn:agentenv:step/v1" EXT_STATE_URI = "urn:agentenv:state/v1" EXT_CHANGES_URI = "urn:agentenv:changes/v1" +EXT_WAKEUPS_URI = "urn:agentenv:wakeups/v1" # Rows per GET /changes page: the default, and the most a client may ask for. CHANGES_PAGE_DEFAULT = 1000 @@ -121,8 +122,9 @@ def sources(names: list[str]) -> str: TRIGGER_UNLESS = "unless" TRIGGER_WEBSITE = "website" TRIGGER_EXECUTOR_HOURS_TIME_ZONE = "executor_hours_time_zone" +TRIGGER_NOTIFY_AGENT = "notify_agent" TRIGGER_FEATURES = (TRIGGER_NAMED_EXECUTORS, TRIGGER_PROVOKING_ROLES, TRIGGER_ACTION_DELAY, TRIGGER_EXECUTOR_HOURS, - TRIGGER_UNLESS, TRIGGER_WEBSITE, TRIGGER_EXECUTOR_HOURS_TIME_ZONE) + TRIGGER_UNLESS, TRIGGER_WEBSITE, TRIGGER_EXECUTOR_HOURS_TIME_ZONE, TRIGGER_NOTIFY_AGENT) GATEWAY_TRAJECTORY_FILE = os.environ.get( "GATEWAY_TRAJECTORY_FILE", @@ -235,4 +237,17 @@ def random_mcp_server_name() -> str: "limit": {"type": "integer", "minimum": 1, "maximum": CHANGES_PAGE_MAX}}}}}}, "required": False, }, + { + "uri": EXT_WAKEUPS_URI, + "description": "Wakeups (urn:agentenv:wakeups/v1): what is due to wake the agent between its turns, on the virtual clock: checks it scheduled for itself (schedule_check, which config turns on) and notifications triggers sent it (the notify_agent action). state says when the next is due; take hands out each due one once (harness control path).", + "params": {"endpoint": "/wakeups/state", "methods": { + "config": {"method": "POST", "endpoint": "/wakeups/config", "request": {"type": "object", "properties": { + "scheduled_checks": {"type": "boolean"}, "max_active": {"type": "integer", "minimum": 1}, + "min_interval": {"type": "string"}}}}, + "state": {"method": "GET", "endpoint": "/wakeups/state", "request": {"type": "object", "properties": {}}}, + "take": {"method": "POST", "endpoint": "/wakeups/take", "request": {"type": "object", "properties": { + "until": {"type": "string"}}}}, + "clear": {"method": "POST", "endpoint": "/wakeups/clear", "request": {"type": "object", "properties": {}}}}}, + "required": False, + }, ] diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index a0d7af367..219e6541f 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -72,6 +72,7 @@ from .get_time import GET_TIME_TOOL_NAME, build_get_time_tool from .rest_proxy import ROLE_HEADER_BYTES, UpstreamResponse, fetch_child_card, forward, rewrite_child_card from .triggers import TriggerEngine, TriggerError +from .wakeups import WAKEUP_TOOL_NAMES, WakeupError, Wakeups, build_wakeup_tools logger = logging.getLogger(__name__) @@ -171,6 +172,7 @@ async def list_website_urls() -> list[str]: ) ) self._clock = Clock() + self._wakeups = Wakeups(lambda: self._clock) self._env_get_time_url_cache: str | None = None self._install_role_filters() @@ -298,6 +300,42 @@ async def clock_state(request: Request) -> Response: st["env_get_time_url"] = await self._env_get_time_url() return self._json_response(st, 200) + async def _wakeup_write(request: Request, method) -> Response: + try: + body = await request.json() + except Exception: + return self._json_response({"ok": False, "error": "Invalid JSON body"}, 400) + try: + return self._json_response(method(body), 200) + except WakeupError as e: + return self._json_response({"ok": False, "error": str(e)}, 400) + + @self._mcp.custom_route("/wakeups/config", methods=["POST"]) + async def wakeups_config(request: Request) -> Response: + taken = sorted(set(WAKEUP_TOOL_NAMES) & set(self._tool_server_urls)) + if taken: + return self._json_response( + {"ok": False, "error": f"{taken} already registered by a backing server"}, 409) + response = await _wakeup_write(request, self._wakeups.configure) + if self._wakeups.enabled: + self._register_wakeup_tools() + else: + self._unregister_wakeup_tools() + return response + + @self._mcp.custom_route("/wakeups/state", methods=["GET"]) + async def wakeups_state(request: Request) -> Response: + return self._json_response(self._wakeups.state(), 200) + + @self._mcp.custom_route("/wakeups/take", methods=["POST"]) + async def wakeups_take(request: Request) -> Response: + return await _wakeup_write(request, self._wakeups.take) + + @self._mcp.custom_route("/wakeups/clear", methods=["POST"]) + async def wakeups_clear(request: Request) -> Response: + self._unregister_wakeup_tools() + return self._json_response(self._wakeups.clear(), 200) + @self._mcp.custom_route("/trajectory", methods=["GET"]) async def env_trajectory(request: Request) -> Response: return self._serve_trajectory() @@ -885,9 +923,11 @@ async def _discover_and_register_tools(self, sessions: dict[str, ClientSession]) raise logger.info(f"Found {len(tools_result.tools)} tool(s) from {server.name}") listed.append((server, tools_result.tools)) - # get_time is left out: a backing server's tool of that name wins, and the clock refuses to arm over it + # get_time and the wakeup tools are left out: a backing server's tool of that name wins, and the clock and + # wakeups refuse to turn on over it sources: dict[str, list[str]] = { - tool.name: ["gateway"] for tool in self._server_tools.get("gateway", []) if tool.name != GET_TIME_TOOL_NAME} + tool.name: ["gateway"] for tool in self._server_tools.get("gateway", []) + if tool.name != GET_TIME_TOOL_NAME and tool.name not in WAKEUP_TOOL_NAMES} for server, tools in listed: for tool in tools: sources.setdefault(tool.name, []).append(server.name) @@ -907,6 +947,8 @@ async def _discover_and_register_tools(self, sessions: dict[str, ClientSession]) self._tools_discovered = True if self._mirror_get_time_descriptor(): logger.info(f"Re-mirrored '{GET_TIME_TOOL_NAME}' descriptor after tool discovery") + if self._mirror_wakeup_descriptors(): + logger.info("Re-mirrored the wakeup tool descriptors after tool discovery") # list_website_urls sits in _server_tools from construction, so a memo may predate the backing tools self._trigger_engine.invalidate_readonly_cache() @@ -1127,6 +1169,46 @@ def _unregister_get_time_tool(self) -> None: self._server_tools.pop("gateway", None) self._trigger_engine.invalidate_readonly_cache() + def _register_wakeup_tools(self) -> None: + """Expose schedule_check and its siblings (idempotent). Called when a task turns scheduled checks on.""" + tools = self._mcp._tool_manager._tools + for tool in build_wakeup_tools(self._wakeups, self._create_arg_model_from_schema): + if tool.name not in self._tool_server_urls: + tools.setdefault(tool.name, tool) + if self._mirror_wakeup_descriptors(): + self._trigger_engine.invalidate_readonly_cache() + + def _mirror_wakeup_descriptors(self) -> bool: + """Mirror the registered wakeup tools into `_server_tools`, as `_mirror_get_time_descriptor` does for + get_time. Returns whether it appended.""" + if not self._tools_discovered or not self._wakeups.enabled: + return False + mirrored = self._server_tools.setdefault("gateway", []) + present = {t.name for t in mirrored} + added = False + for name in WAKEUP_TOOL_NAMES: + tool = self._mcp._tool_manager._tools.get(name) + if tool is None or name in present or name in self._tool_server_urls: + continue + mirrored.append(MCPTool(name=tool.name, description=tool.description, + inputSchema=tool.parameters, annotations=tool.annotations)) + added = True + return added + + def _unregister_wakeup_tools(self) -> None: + """Remove them when a task turns scheduled checks off or clears (idempotent); a backing server's are left.""" + tools = self._mcp._tool_manager._tools + for name in WAKEUP_TOOL_NAMES: + if name not in self._tool_server_urls: + tools.pop(name, None) + remaining = [t for t in self._server_tools.get("gateway", []) if t.name not in WAKEUP_TOOL_NAMES + or t.name in self._tool_server_urls] + if remaining: + self._server_tools["gateway"] = remaining + else: + self._server_tools.pop("gateway", None) + self._trigger_engine.invalidate_readonly_cache() + def _create_and_register_proxy_tool(self, server_url: str, tool: MCPTool) -> None: self._mcp._tool_manager._tools[tool.name] = self._proxy_tool(server_url, tool) diff --git a/src/agent_env/env/gateway/triggers.py b/src/agent_env/env/gateway/triggers.py index 8d4eb8817..bb2daa230 100644 --- a/src/agent_env/env/gateway/triggers.py +++ b/src/agent_env/env/gateway/triggers.py @@ -23,6 +23,7 @@ from mcp.types import CallToolResult, TextContent from .clock import ClockError, _advance, _iso, _parse_duration, _parse_rfc3339, _zone +from .wakeups import validate_channel from .caller import Caller from .constants import DEFAULT_ROLE, TRIGGER_IN_FLIGHT_STATUSES, TRIGGER_STATUSES @@ -34,7 +35,7 @@ _CTX_ROOTS = {"action": ("args", "result"), "website": ("args",)} # RFC 9110's safe methods: a request using one asks for no change, so it prompts no state re-evaluation. _SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE"}) -_ACTION_TYPES = ("nl", "tool", "permission") +_ACTION_TYPES = ("nl", "tool", "permission", "notify_agent") _TERMINAL_A2A_STATES = ("completed", "failed", "canceled", "rejected") # Bounds the per-evaluation catch-up burst; the backlog drains on later evaluations. _MAX_CATCHUP = 1000 @@ -741,6 +742,8 @@ def _normalize_action(action: Any, w: str, roots: tuple[str, ...], executors: di raise TriggerError(f"{w}: 'as' is only valid on tool actions") if "executor" in action and atype != "nl": raise TriggerError(f"{w}: 'executor' is only valid on nl actions") + if "verify" in action and atype == "notify_agent": + raise TriggerError(f"{w}: a notify_agent action has nothing to verify") if action.get("delay") is None: action.pop("delay", None) else: @@ -760,6 +763,14 @@ def _normalize_action(action: Any, w: str, roots: tuple[str, ...], executors: di elif not isinstance(named, str) or named not in executors: raise TriggerError(f"{w}: executor {named!r} is not a registered executor name " f"(registered: {sorted(n for n in executors if n is not None)})") + elif atype == "notify_agent": + try: + validate_channel(action.get("channel"), w) + except ValueError as e: + raise TriggerError(str(e)) + if not isinstance(action.get("text"), str) or not action["text"].strip(): + raise TriggerError(f"{w}: text must be the notification's non-empty text") + _validate_placeholders(action["text"], f"{w}.text", roots) elif atype == "tool": if not isinstance(action.get("tool"), str) or not action["tool"]: raise TriggerError(f"{w}: tool must be a non-empty string") @@ -1317,6 +1328,15 @@ async def _run_action(self, tid: str, trig: dict, index: int, action: dict, ctx: tool=action["tool"], args=_echo_args(args), **acting) return False detail = {"tool": action["tool"], "args": _echo_args(args), **acting} + elif action["type"] == "notify_agent": + try: + text = _template(action["text"], {}, ctx) + except TemplateError as e: + self._emit("action_failed", tid, action_index=index, detail=str(e)[:300]) + return False + note = self._gw._wakeups.notify(action["channel"], str(text), trigger_id=tid) + # the text may carry what a person wrote, so the log keeps its size, not its words + detail = {"notify_agent": action["channel"], "notification_id": note["id"], "text_len": len(str(text))} else: if not await self._run_nl(tid, trig, index, action, ctx): return False diff --git a/src/agent_env/env/gateway/wakeups.py b/src/agent_env/env/gateway/wakeups.py new file mode 100644 index 000000000..7872c137b --- /dev/null +++ b/src/agent_env/env/gateway/wakeups.py @@ -0,0 +1,286 @@ +"""What wakes an agent between its turns (urn:agentenv:wakeups/v1). + +Two sources, both on the gateway's virtual clock: +- checks the agent schedules for itself with the schedule_check / list_scheduled_checks / cancel_scheduled_check + tools, which exist only once a task turns them on (`configure`); +- notifications a task's triggers send it with the `notify_agent` action, through a channel the task declares. + +Nothing here wakes anyone: a harness (a post-turn hook) reads `state` to learn when the next one is due, moves the +clock there, and `take`s what has come due to put it in front of the agent. +""" +from __future__ import annotations + +import threading +from datetime import datetime, timedelta +from typing import Any, Callable + +from mcp.server.fastmcp.tools import Tool as FastMCPTool +from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata +from mcp.types import CallToolResult, TextContent, ToolAnnotations + +from .clock import Clock, ClockError, _iso, _parse_duration, _parse_rfc3339 + +SCHEDULE_CHECK_TOOL_NAME = "schedule_check" +LIST_SCHEDULED_CHECKS_TOOL_NAME = "list_scheduled_checks" +CANCEL_SCHEDULED_CHECK_TOOL_NAME = "cancel_scheduled_check" +WAKEUP_TOOL_NAMES = (SCHEDULE_CHECK_TOOL_NAME, LIST_SCHEDULED_CHECKS_TOOL_NAME, CANCEL_SCHEDULED_CHECK_TOOL_NAME) + +DEFAULT_MAX_ACTIVE = 10 +DEFAULT_MIN_INTERVAL = "PT1H" +_MAX_NOTE_CHARS = 2000 +_MAX_CHANNEL_CHARS = 60 + +SCHEDULE_CHECK_DESCRIPTION = ( + "Schedule a check-in for yourself. You only run when you are messaged, so to look at something again later " + "(a reply you are waiting for, a listing to watch, a deadline), schedule a check: when it comes due you are " + "woken with your note, so write the note so it tells you what to do. Give `at` (a time, or a duration from now " + "such as PT6H) and/or `every` (a duration such as P1D, to repeat), with an optional `until`. Times are in this " + "environment's clock (see get_time); a time without a UTC offset is refused." +) +_SCHEDULE_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "note": {"type": "string", "description": "What to do when the check comes due."}, + "at": {"type": "string", "description": "When: an RFC 3339 time with its offset (2027-02-02T18:00:00-08:00), " + "or an ISO 8601 duration from now (PT6H)."}, + "every": {"type": "string", "description": "Repeat this often: an ISO 8601 duration (PT6H, P1D)."}, + "until": {"type": "string", "description": "Stop repeating after this time (RFC 3339) or duration from now."}, + }, + "required": ["note"], + "additionalProperties": False, +} +_LIST_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}, "additionalProperties": False} +_CANCEL_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": {"check_id": {"type": "string"}}, + "required": ["check_id"], + "additionalProperties": False, +} + + +class WakeupError(ValueError): + """A bad configuration or tool argument (a 400 on the harness routes, a tool error for the agent).""" + + +def _when(text: Any, now: datetime, field: str) -> datetime: + """An RFC 3339 instant, or an ISO 8601 duration counted from now.""" + if not isinstance(text, str) or not text: + raise WakeupError(f"{field} must be an RFC 3339 time or an ISO 8601 duration") + try: + return now + _parse_duration(text) if text.startswith("P") else _parse_rfc3339(text) + except ClockError as e: + raise WakeupError(f"{field}: {e}") + + +class Wakeups: + """The gateway's record of scheduled checks and agent notifications; threadsafe. `clock` returns the gateway's + current Clock, read at each use so a swapped clock (tests) is followed.""" + + def __init__(self, clock: Callable[[], Clock]): + self._clock_of = clock + self._lock = threading.Lock() + self._enabled = False + self._max_active = DEFAULT_MAX_ACTIVE + self._min_interval = _parse_duration(DEFAULT_MIN_INTERVAL) + self._checks: dict[str, dict] = {} + self._notifications: list[dict] = [] + self._seq = 0 + + @property + def enabled(self) -> bool: + return self._enabled + + def configure(self, body: Any) -> dict: + """Turn scheduled checks on or off, with their limits. Notifications need no switch.""" + if not isinstance(body, dict): + raise WakeupError("body must be a JSON object") + enabled = body.get("scheduled_checks", True) + if not isinstance(enabled, bool): + raise WakeupError("scheduled_checks must be a boolean") + max_active = body.get("max_active", DEFAULT_MAX_ACTIVE) + if isinstance(max_active, bool) or not isinstance(max_active, int) or max_active < 1: + raise WakeupError("max_active must be a positive integer") + try: + min_interval = _parse_duration(body.get("min_interval", DEFAULT_MIN_INTERVAL)) + except ClockError as e: + raise WakeupError(f"min_interval: {e}") + if min_interval <= timedelta(0): + raise WakeupError("min_interval must be a positive duration") + with self._lock: + self._enabled, self._max_active, self._min_interval = enabled, max_active, min_interval + return self.state() + + def clear(self) -> dict: + with self._lock: + self._enabled = False + self._checks.clear() + self._notifications.clear() + return {"ok": True} + + @property + def _clock(self) -> Clock: + return self._clock_of() + + def _now(self) -> datetime: + now = self._clock.now() + if now is None: + raise WakeupError("this environment's clock is not set, so checks can't be scheduled") + return now + + def _next_id(self, prefix: str) -> str: + self._seq += 1 + return f"{prefix}-{self._seq}" + + # ---- agent tools ---- + + def schedule(self, note: Any, at: Any = None, every: Any = None, until: Any = None) -> dict: + if not self._enabled: + raise WakeupError("scheduled checks are not available here") + if not isinstance(note, str) or not note.strip(): + raise WakeupError("note must say what to do when the check comes due") + if len(note) > _MAX_NOTE_CHARS: + raise WakeupError(f"note must be at most {_MAX_NOTE_CHARS} characters") + now = self._now() + step = None + if every is not None: + try: + step = _parse_duration(every) + except ClockError as e: + raise WakeupError(f"every: {e}") + if step < self._min_interval: + raise WakeupError(f"every must be at least {self._min_interval} apart") + first = _when(at, now, "at") if at is not None else (now + step if step else None) + if first is None: + raise WakeupError("give at (when) or every (how often)") + if first <= now: + raise WakeupError(f"at {_iso(first)} has already passed; it is now {_iso(now)}") + end = _when(until, now, "until") if until is not None else None + if end is not None and step is None: + raise WakeupError("until only applies with every") + if end is not None and end < first: + raise WakeupError("until is before the first check") + with self._lock: + if sum(c["status"] == "active" for c in self._checks.values()) >= self._max_active: + raise WakeupError(f"you already have {self._max_active} scheduled checks; cancel one first") + cid = self._next_id("check") + self._checks[cid] = {"check_id": cid, "note": note, "next": first, "every": step, "until": end, + "created_at": now, "status": "active", "fired": 0} + return self._public(self._checks[cid]) + + def list(self) -> dict: + with self._lock: + return {"checks": [self._public(c) for c in self._checks.values() if c["status"] == "active"]} + + def cancel(self, check_id: Any) -> dict: + with self._lock: + check = self._checks.get(check_id) if isinstance(check_id, str) else None + if check is None or check["status"] != "active": + raise WakeupError(f"no active scheduled check {check_id!r}") + check["status"] = "cancelled" + return {"cancelled": check_id} + + @staticmethod + def _public(check: dict) -> dict: + out = {"check_id": check["check_id"], "note": check["note"], "next": _iso(check["next"])} + if check["every"]: + out["every"] = str(check["every"]) + if check["until"]: + out["until"] = _iso(check["until"]) + return out + + # ---- trigger action ---- + + def notify(self, channel: str, text: str, trigger_id: str | None = None) -> dict: + """Queue a notification for the agent at the current virtual time (wall time if the clock is unset).""" + now = self._clock.now() + with self._lock: + item = {"id": self._next_id("note"), "channel": channel, "text": text, "trigger_id": trigger_id, + "at": _iso(now) if now else None, "status": "pending"} + self._notifications.append(item) + return dict(item) + + # ---- harness ---- + + def state(self) -> dict: + with self._lock: + active = [c for c in self._checks.values() if c["status"] == "active"] + pending = [dict(n) for n in self._notifications if n["status"] == "pending"] + nxt = min((c["next"] for c in active), default=None) + return {"scheduled_checks": self._enabled, "max_active": self._max_active, + "min_interval": str(self._min_interval), "next_check_at": _iso(nxt) if nxt else None, + "checks": [{**self._public(c), "status": c["status"], "fired": c["fired"]} + for c in self._checks.values()], + "pending_notifications": pending} + + def take(self, body: Any = None) -> dict: + """Everything due by `until` (default: now), oldest first; each is handed out once. A repeating check moves to + its next time after `until` and says how many times it came due, so a long jump wakes the agent once.""" + body = body if isinstance(body, dict) else {} + now = self._clock.now() + until = _when(body["until"], now, "until") if body.get("until") is not None and now else now + items: list[dict] = [] + with self._lock: + for check in self._checks.values(): + if check["status"] != "active" or until is None or check["next"] > until: + continue + due, times = check["next"], 0 + while check["status"] == "active" and check["next"] <= until: + times += 1 + if check["every"] and (check["until"] is None or check["next"] + check["every"] <= check["until"]): + check["next"] = check["next"] + check["every"] + else: + check["status"] = "done" + check["fired"] += 1 + items.append({"kind": "scheduled_check", "check_id": check["check_id"], "note": check["note"], + "at": _iso(due), "times": times}) + for note in self._notifications: + if note["status"] == "pending" and (until is None or note["at"] is None or note["at"] <= _iso(until)): + note["status"] = "delivered" + items.append({"kind": "notification", "id": note["id"], "channel": note["channel"], + "text": note["text"], "at": note["at"], "trigger_id": note["trigger_id"]}) + items.sort(key=lambda i: i["at"] or "") + return {"items": items, "until": _iso(until) if until else None} + + +def _tool_error(text: str) -> CallToolResult: + return CallToolResult(isError=True, content=[TextContent(type="text", text=text)]) + + +def build_wakeup_tools(wakeups: Wakeups, make_arg_model: Callable[[str, dict], Any]) -> list[FastMCPTool]: + """The agent-facing tools, built by hand like get_time so a backing server's tool can't silently shadow them.""" + + async def schedule_check(note: str, at: str | None = None, every: str | None = None, + until: str | None = None) -> dict | CallToolResult: + try: + return wakeups.schedule(note, at=at, every=every, until=until) + except WakeupError as e: + return _tool_error(str(e)) + + async def list_scheduled_checks() -> dict: + return wakeups.list() + + async def cancel_scheduled_check(check_id: str) -> dict | CallToolResult: + try: + return wakeups.cancel(check_id) + except WakeupError as e: + return _tool_error(str(e)) + + specs = [ + (SCHEDULE_CHECK_TOOL_NAME, schedule_check, SCHEDULE_CHECK_DESCRIPTION, _SCHEDULE_SCHEMA, False), + (LIST_SCHEDULED_CHECKS_TOOL_NAME, list_scheduled_checks, "Your scheduled checks that are still to come.", + _LIST_SCHEMA, True), + (CANCEL_SCHEDULED_CHECK_TOOL_NAME, cancel_scheduled_check, "Cancel one of your scheduled checks.", + _CANCEL_SCHEMA, False), + ] + return [FastMCPTool(fn=fn, name=name, description=description, parameters=schema, + fn_metadata=FuncMetadata(arg_model=make_arg_model(name, schema)), is_async=True, + annotations=ToolAnnotations(readOnlyHint=read_only)) + for name, fn, description, schema, read_only in specs] + + +def validate_channel(channel: Any, where: str) -> None: + """A notify_agent channel: the short source label the agent sees, e.g. 'Gmail' or 'Curbfinch alert'.""" + if not isinstance(channel, str) or not channel.strip(): + raise ValueError(f"{where}: channel must be a non-empty string, e.g. 'Gmail' or 'Curbfinch alert'") + if len(channel) > _MAX_CHANNEL_CHARS: + raise ValueError(f"{where}: channel must be at most {_MAX_CHANNEL_CHARS} characters") diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 5090f6efe..937bfc9de 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -147,7 +147,8 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s UNAFFECTED = frozenset({ # They name no agent. - "aggregate_verifiers", "apply_server_config", "build_mcp_cli", "combine_universe_verdicts", "deploy_env", + "aggregate_verifiers", "apply_server_config", "build_mcp_cli", "combine_universe_verdicts", "configure_wakeups", + "deploy_env", "deploy_sandbox", "env_outcome_verifier", "modify_env_tool_access", "reset_env", "review", "run_container_unit_tests_verifier", "run_docker_container", "snapshot_env", "sync_env_clock", "validation_gate_aggregator", "verify_env_card", "verify_env_core_protocol", "verify_mcp_tool_schema", diff --git a/src/agent_env/task_step/registry.py b/src/agent_env/task_step/registry.py index 76d1934c0..8180e4be3 100644 --- a/src/agent_env/task_step/registry.py +++ b/src/agent_env/task_step/registry.py @@ -42,6 +42,7 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: from agent_env.task_step.task_steps.register_env_triggers import RegisterEnvTriggersStep from agent_env.task_step.task_steps.register_agent_triggers import RegisterAgentTriggersStep from agent_env.task_step.task_steps.sync_env_clock import SyncEnvClockTaskStep + from agent_env.task_step.task_steps.configure_wakeups import ConfigureWakeupsTaskStep from agent_env.task_step.task_steps.install_agent import InstallAgentTaskStep from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep from agent_env.task_step.task_steps.verifiers.run_container_unit_tests_verifier import RunContainerUnitTestsVerifierTaskStep @@ -94,6 +95,7 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: _get_type(RegisterEnvTriggersStep): RegisterEnvTriggersStep, _get_type(RegisterAgentTriggersStep): RegisterAgentTriggersStep, _get_type(SyncEnvClockTaskStep): SyncEnvClockTaskStep, + _get_type(ConfigureWakeupsTaskStep): ConfigureWakeupsTaskStep, _get_type(InstallAgentTaskStep): InstallAgentTaskStep, _get_type(RubricsVerifierTaskStep): RubricsVerifierTaskStep, _get_type(RunContainerUnitTestsVerifierTaskStep): RunContainerUnitTestsVerifierTaskStep, diff --git a/src/agent_env/task_step/task_steps/configure_wakeups.py b/src/agent_env/task_step/task_steps/configure_wakeups.py new file mode 100644 index 000000000..0700c37f2 --- /dev/null +++ b/src/agent_env/task_step/task_steps/configure_wakeups.py @@ -0,0 +1,90 @@ +"""Task step that turns on the checks an agent can schedule for itself on a deployed env's gateway (wakeups/v1). + +With it the agent gets schedule_check, list_scheduled_checks and cancel_scheduled_check, on the env's virtual clock. +Place it after sync_env_clock (a check needs the clock) and before prompt_agent (the agent needs the tools on its +first turn). A post-turn hook reads the gateway's wakeups state to wake the agent when a check comes due. +""" + +from __future__ import annotations + +import logging +from typing import ClassVar, Optional + +import httpx + +from agent_env.env.gateway.constants import EXT_WAKEUPS_URI +from agent_env.env.gateway.wakeups import DEFAULT_MAX_ACTIVE, DEFAULT_MIN_INTERVAL +from agent_env.entity.refs import EntityRef +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep, TaskStepDependency + +logger = logging.getLogger(__name__) + + +class ConfigureWakeupsTaskStep(TaskStep): + type: ClassVar[str] = "configure_wakeups" + entity_refs = (EntityRef.env("env_id"),) + + def __init__( + self, + id: str, + version: Optional[int], + env_id: str, + scheduled_checks: bool = True, + max_active: int = DEFAULT_MAX_ACTIVE, + min_interval: str = DEFAULT_MIN_INTERVAL, + depends_on: Optional[list[TaskStepDependency]] = None, + fail_task_on_error: bool = True, + timeout_seconds: int = 30, + ): + """max_active caps the agent's checks still to come; min_interval is the shortest `every` it may ask for.""" + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + if not isinstance(scheduled_checks, bool): + raise ValueError("scheduled_checks must be a boolean") + if isinstance(max_active, bool) or not isinstance(max_active, int) or max_active < 1: + raise ValueError("max_active must be a positive integer") + self.env_id = env_id + self.scheduled_checks = scheduled_checks + self.max_active = max_active + self.min_interval = min_interval + self.timeout_seconds = timeout_seconds + + def to_dict(self) -> dict: + base = super().to_dict() + base["env_id"] = self.env_id + base["scheduled_checks"] = self.scheduled_checks + base["max_active"] = self.max_active + base["min_interval"] = self.min_interval + base["timeout_seconds"] = self.timeout_seconds + return base + + @classmethod + def from_dict(cls, data: dict) -> ConfigureWakeupsTaskStep: + return cls( + **cls._base_from_dict(data), + env_id=data["env_id"], + scheduled_checks=data.get("scheduled_checks", True), + max_active=data.get("max_active", DEFAULT_MAX_ACTIVE), + min_interval=data.get("min_interval", DEFAULT_MIN_INTERVAL), + timeout_seconds=data.get("timeout_seconds", 30), + ) + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) + if deployed is None: + raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + if not deployed.supports(EXT_WAKEUPS_URI, "config"): + raise RuntimeError(f"Env '{self.env_id}' runs a gateway without {EXT_WAKEUPS_URI}; redeploy it on a gateway " + "image that has it.") + body = {"scheduled_checks": self.scheduled_checks, "max_active": self.max_active, + "min_interval": self.min_interval} + try: + state = await deployed.invoke(EXT_WAKEUPS_URI, "config", body, timeout=self.timeout_seconds) + except httpx.HTTPStatusError as e: + raise RuntimeError(f"wakeups config failed (HTTP {e.response.status_code}): {e.response.text}") from e + logger.info(f"configure_wakeups env={self.env_id}: {body}") + context.metadata.setdefault("wakeup_configurations", []).append({ + "step_id": self.id, "env_id": self.env_id, **body, + "applied": {k: state.get(k) for k in ("scheduled_checks", "max_active", "min_interval")}, + }) + return context diff --git a/src/agent_env/task_step/task_steps/register_env_triggers.py b/src/agent_env/task_step/task_steps/register_env_triggers.py index cf651595a..b36970ded 100644 --- a/src/agent_env/task_step/task_steps/register_env_triggers.py +++ b/src/agent_env/task_step/task_steps/register_env_triggers.py @@ -17,6 +17,7 @@ TRIGGER_EXECUTOR_HOURS, TRIGGER_EXECUTOR_HOURS_TIME_ZONE, TRIGGER_NAMED_EXECUTORS, + TRIGGER_NOTIFY_AGENT, TRIGGER_PROVOKING_ROLES, TRIGGER_UNLESS, TRIGGER_WEBSITE, @@ -151,7 +152,9 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: TRIGGER_UNLESS: _when_has(self.triggers, "unless"), TRIGGER_WEBSITE: any(isinstance(t.get("when"), dict) and t["when"].get("type") == "website" for t in self.triggers), - TRIGGER_EXECUTOR_HOURS_TIME_ZONE: _hours_in_a_time_zone(self.executor_hours)} + TRIGGER_EXECUTOR_HOURS_TIME_ZONE: _hours_in_a_time_zone(self.executor_hours), + TRIGGER_NOTIFY_AGENT: any(a.get("type") == "notify_agent" for t in self.triggers + for a in t.get("actions") or [] if isinstance(a, dict))} missing = [feature for feature, needed in used.items() if needed and feature not in _trigger_features(deployed)] if missing: raise RuntimeError( diff --git a/tst/unit/env/gateway/gateway_card_test.py b/tst/unit/env/gateway/gateway_card_test.py index 05b06ce69..704cc02c2 100644 --- a/tst/unit/env/gateway/gateway_card_test.py +++ b/tst/unit/env/gateway/gateway_card_test.py @@ -201,12 +201,12 @@ def test_rewrite_child_card_drops_the_childs_interfaces(): assert _gateway({"mcp-slack": "http://slack:18765"})._rewrite_child_card("mcp-slack", child)["additionalInterfaces"] == [] -def test_gateway_card_lists_its_eight_extensions(): +def test_gateway_card_lists_its_nine_extensions(): assert _gateway({})._gateway_extensions() is gateway_constants.GATEWAY_EXTENSIONS assert [e["uri"] for e in _gateway({})._gateway_extensions()] == [ "urn:agentenv:disable-tool/v1", "urn:agentenv:enable-tool/v1", "urn:agentenv:triggers/v1", "urn:agentenv:clock/v1", "urn:agentenv:trajectory/v1", "urn:agentenv:step/v1", "urn:agentenv:state/v1", - "urn:agentenv:changes/v1", + "urn:agentenv:changes/v1", "urn:agentenv:wakeups/v1", ] diff --git a/tst/unit/env/gateway/wakeups_test.py b/tst/unit/env/gateway/wakeups_test.py new file mode 100644 index 000000000..85f2c35b8 --- /dev/null +++ b/tst/unit/env/gateway/wakeups_test.py @@ -0,0 +1,289 @@ +"""Unit tests for what wakes an agent between turns (urn:agentenv:wakeups/v1): scheduled checks and notify_agent.""" +from __future__ import annotations + +import asyncio +import json + +import httpx +import pytest +from mcp.types import CallToolResult, TextContent +from mcp.types import Tool as MCPTool + +from agent_env.env.gateway import gateway as gw_mod +from agent_env.env.gateway.clock import Clock +from agent_env.env.gateway.gateway import Gateway +from agent_env.env.gateway.triggers import TriggerEngine, TriggerError +from agent_env.env.gateway.wakeups import WAKEUP_TOOL_NAMES, WakeupError, Wakeups + +T0 = "2027-02-01T15:00:00Z" # Mon 7:00 AM Pacific + + +class FakeNow: + def __init__(self, start: float = 1000.0): + self.t = start + + def __call__(self) -> float: + return self.t + + def tick(self, seconds: float) -> None: + self.t += seconds + + +def _frozen_clock(at: str = T0) -> Clock: + clock = Clock(now_fn=FakeNow()) + clock.set_time(at, 0) + return clock + + +def _wakeups(clock: Clock | None = None, **config) -> Wakeups: + clock = clock or _frozen_clock() + w = Wakeups(lambda: clock) + w.configure({"scheduled_checks": True, **config}) + return w + + +# ---- Wakeups ---- + +def test_scheduled_checks_are_off_until_a_task_turns_them_on(): + w = Wakeups(lambda: _frozen_clock()) + assert w.state()["scheduled_checks"] is False + with pytest.raises(WakeupError, match="not available"): + w.schedule("look again", at="PT6H") + + +def test_a_check_takes_a_duration_from_now_or_a_time_with_its_offset(): + w = _wakeups() + assert w.schedule("in six hours", at="PT6H")["next"] == "2027-02-01T21:00:00Z" + assert w.schedule("tonight", at="2027-02-01T18:00:00-08:00")["next"] == "2027-02-02T02:00:00Z" + assert [c["note"] for c in w.list()["checks"]] == ["in six hours", "tonight"] + + +@pytest.mark.parametrize("kwargs, error", [ + ({"note": "x"}, "give at"), + ({"note": " ", "at": "PT1H"}, "note must say"), + ({"note": "x", "at": "2027-02-01T06:00:00-08:00"}, "already passed"), + ({"note": "x", "at": "2027-02-01T18:00:00"}, "RFC 3339|RFC3339"), + ({"note": "x", "every": "PT10M"}, "at least"), + ({"note": "x", "at": "PT1H", "until": "P1D"}, "until only applies"), + ({"note": "x", "every": "PT6H", "until": "PT1H"}, "before the first"), +]) +def test_bad_checks_are_refused_with_the_reason(kwargs, error): + with pytest.raises(WakeupError, match=error): + _wakeups().schedule(**kwargs) + + +def test_the_active_cap_refuses_one_more_until_one_is_cancelled(): + w = _wakeups(max_active=2) + first = w.schedule("a", at="PT1H")["check_id"] + w.schedule("b", at="PT2H") + with pytest.raises(WakeupError, match="already have 2"): + w.schedule("c", at="PT3H") + w.cancel(first) + w.schedule("c", at="PT3H") + with pytest.raises(WakeupError, match="no active scheduled check"): + w.cancel(first) + + +def test_take_hands_out_each_due_check_once_and_a_one_shot_is_done(): + clock = _frozen_clock() + w = _wakeups(clock) + w.schedule("look at 9", at="PT2H") + w.schedule("look at noon", at="PT5H") + assert w.state()["next_check_at"] == "2027-02-01T17:00:00Z" + clock.set_time("2027-02-01T17:00:00Z", 0) + taken = w.take({})["items"] + assert [(i["kind"], i["note"], i["at"]) for i in taken] == [("scheduled_check", "look at 9", "2027-02-01T17:00:00Z")] + assert w.take({})["items"] == [] + assert [c["note"] for c in w.list()["checks"]] == ["look at noon"] + assert w.state()["next_check_at"] == "2027-02-01T20:00:00Z" + + +def test_a_repeating_check_moves_past_a_long_jump_and_says_how_often_it_came_due(): + clock = _frozen_clock() + w = _wakeups(clock) + w.schedule("check rugs", at="PT5H", every="PT6H", until="P4D") + clock.set_time("2027-02-02T15:00:00Z", 0) # a day later: due at 20:00, 02:00, 08:00, 14:00 + [item] = w.take({})["items"] + assert (item["at"], item["times"]) == ("2027-02-01T20:00:00Z", 4) + assert w.list()["checks"][0]["next"] == "2027-02-02T20:00:00Z" + + +def test_a_repeating_check_ends_at_until(): + clock = _frozen_clock() + w = _wakeups(clock) + w.schedule("check", every="P1D", until="P2D") + clock.set_time("2027-02-05T15:00:00Z", 0) + assert w.take({})["items"][0]["times"] == 2 + assert w.list()["checks"] == [] and w.state()["next_check_at"] is None + + +def test_take_until_looks_ahead_without_moving_the_clock(): + w = _wakeups() + w.schedule("later", at="PT3H") + assert w.take({"until": "PT1H"})["items"] == [] + assert [i["note"] for i in w.take({"until": "PT3H"})["items"]] == ["later"] + + +def test_notifications_queue_at_the_virtual_time_and_are_taken_once(): + clock = _frozen_clock() + w = Wakeups(lambda: clock) # notifications need no switch + w.notify("Messages", "Jamie: did the rug thing get sorted?", trigger_id="t1") + pending = w.state()["pending_notifications"] + assert [(n["channel"], n["at"], n["trigger_id"]) for n in pending] == [("Messages", T0, "t1")] + [item] = w.take({})["items"] + assert (item["kind"], item["text"]) == ("notification", "Jamie: did the rug thing get sorted?") + assert w.take({})["items"] == [] and w.state()["pending_notifications"] == [] + + +def test_clear_turns_checks_off_and_forgets_everything(): + w = _wakeups() + w.schedule("x", at="PT1H") + w.notify("Gmail", "hello") + w.clear() + st = w.state() + assert (st["scheduled_checks"], st["checks"], st["pending_notifications"]) == (False, [], []) + + +@pytest.mark.parametrize("body, error", [ + ({"scheduled_checks": "yes"}, "boolean"), + ({"max_active": 0}, "positive integer"), + ({"min_interval": "1h"}, "min_interval"), +]) +def test_bad_config_is_refused(body, error): + with pytest.raises(WakeupError, match=error): + Wakeups(lambda: _frozen_clock()).configure(body) + + +# ---- the gateway ---- + +@pytest.fixture +def gateway(tmp_path, monkeypatch) -> Gateway: + monkeypatch.setattr(gw_mod, "GATEWAY_TRAJECTORY_FILE", str(tmp_path / "trajectory.jsonl")) + gw = Gateway(host="127.0.0.1", port=0, server_name="AgentEnvGateway", internal_mcp_servers=[]) + gw._clock = _frozen_clock() + gw._server_tools["mcp-svc"] = [MCPTool(name="svc_thing", description="", inputSchema={"type": "object"})] + gw._tools_discovered = True + yield gw + gw._close_trajectory_file() + + +def _request(gw: Gateway, method: str, path: str, body: dict | None = None) -> httpx.Response: + async def _do() -> httpx.Response: + transport = httpx.ASGITransport(app=gw._mcp.streamable_http_app()) + async with httpx.AsyncClient(transport=transport, base_url="http://gateway") as client: + return await client.request(method, path, json=body) + + return asyncio.run(_do()) + + +def _gateway_tool_names(gw: Gateway) -> set[str]: + return {t.name for t in gw._server_tools.get("gateway", [])} + + +def test_the_tools_appear_only_once_a_task_turns_them_on_and_go_on_clear(gateway): + tools = gateway._mcp._tool_manager._tools + assert not set(WAKEUP_TOOL_NAMES) & set(tools) + r = _request(gateway, "POST", "/wakeups/config", {"max_active": 5}) + assert r.status_code == 200 and r.json()["max_active"] == 5 + assert set(WAKEUP_TOOL_NAMES) <= set(tools) and set(WAKEUP_TOOL_NAMES) <= _gateway_tool_names(gateway) + assert _request(gateway, "POST", "/wakeups/clear", {}).status_code == 200 + assert not set(WAKEUP_TOOL_NAMES) & set(tools) and not set(WAKEUP_TOOL_NAMES) & _gateway_tool_names(gateway) + + +def test_turning_checks_off_removes_the_tools(gateway): + _request(gateway, "POST", "/wakeups/config", {}) + _request(gateway, "POST", "/wakeups/config", {"scheduled_checks": False}) + assert not set(WAKEUP_TOOL_NAMES) & set(gateway._mcp._tool_manager._tools) + + +def test_config_is_refused_over_a_backing_server_tool_of_the_same_name(gateway): + gateway._tool_server_urls["schedule_check"] = "http://svc:1/mcp" + r = _request(gateway, "POST", "/wakeups/config", {}) + assert r.status_code == 409 and "schedule_check" in r.json()["error"] + + +def test_a_bad_config_is_a_400(gateway): + assert _request(gateway, "POST", "/wakeups/config", {"max_active": -1}).status_code == 400 + + +def test_the_agent_schedules_through_the_tool_and_the_harness_takes_it(gateway): + _request(gateway, "POST", "/wakeups/config", {}) + tool = gateway._mcp._tool_manager._tools["schedule_check"] + made = asyncio.run(tool.fn(note="check rugs", at="PT2H")) + assert made["next"] == "2027-02-01T17:00:00Z" + refused = asyncio.run(tool.fn(note="check", every="PT5M")) + assert isinstance(refused, CallToolResult) and refused.isError + state = _request(gateway, "GET", "/wakeups/state").json() + assert state["next_check_at"] == "2027-02-01T17:00:00Z" + gateway._clock.set_time("2027-02-01T17:00:00Z", 0) + items = _request(gateway, "POST", "/wakeups/take", {}).json()["items"] + assert [i["note"] for i in items] == ["check rugs"] + + +def test_list_is_read_only_so_it_does_not_re_check_state_triggers(gateway): + _request(gateway, "POST", "/wakeups/config", {}) + annotations = {t.name: t.annotations for t in gateway._server_tools["gateway"]} + assert annotations["list_scheduled_checks"].readOnlyHint is True + assert annotations["schedule_check"].readOnlyHint is False + + +# ---- notify_agent ---- + +class _FakeGateway: + def __init__(self) -> None: + self.events: list[dict] = [] + self._role_rules_lock = asyncio.Lock() + self._server_tools: dict = {} + self._tool_server_urls: dict = {} + self._clock = _frozen_clock() + self._wakeups = Wakeups(lambda: self._clock) + self.TOOL_CALL_TIMEOUT_S = 30 + + def _query_changelog_id(self): + return 1 + + async def _log_event(self, event): + self.events.append(dict(event)) + return f"event_{len(self.events)}" + + async def _ensure_tools_discovered(self): + pass + + +def _ok(text: str = "{}") -> CallToolResult: + return CallToolResult(content=[TextContent(type="text", text=text)], isError=False) + + +def _notify_trigger(**action) -> dict: + return {"id": "jamie-texts", "when": {"type": "action", "tool": "messaging_send_message"}, + "actions": [{"type": "notify_agent", "channel": "Messages", "text": "Jamie: ok", **action}]} + + +@pytest.mark.asyncio +async def test_notify_agent_queues_a_templated_notification(): + gw = _FakeGateway() + engine = TriggerEngine(gw) + engine.register({"watch_roles": ["default"], + "triggers": [_notify_trigger(text="${args.to}: got your text about ${args.subject}")]}) + engine.on_tool_call("default", "messaging_send_message", {"to": "Jamie", "subject": "the rug"}, _ok()) + await asyncio.sleep(0.2) + [note] = gw._wakeups.state()["pending_notifications"] + assert (note["channel"], note["text"], note["at"]) == ("Messages", "Jamie: got your text about the rug", T0) + ok = [e for e in engine.state()["events"] if e["kind"] == "action_ok"] + assert ok and ok[0]["detail"]["notify_agent"] == "Messages" and "got your text" not in json.dumps(ok[0]) + + +@pytest.mark.parametrize("action, error", [ + ({"channel": ""}, "channel"), + ({"text": " "}, "text"), + ({"verify": {"tool": "x", "predicate": {"exists": True}}}, "nothing to verify"), + ({"text": "${result.id}", "channel": "Gmail"}, None), +]) +def test_notify_agent_is_validated_at_registration(action, error): + engine = TriggerEngine(_FakeGateway()) + body = {"watch_roles": ["default"], "triggers": [_notify_trigger(**action)]} + if error is None: + engine.register(body) + else: + with pytest.raises(TriggerError, match=error): + engine.register(body) diff --git a/tst/unit/task_step/test_configure_wakeups.py b/tst/unit/task_step/test_configure_wakeups.py new file mode 100644 index 000000000..ff04738f1 --- /dev/null +++ b/tst/unit/task_step/test_configure_wakeups.py @@ -0,0 +1,79 @@ +"""ConfigureWakeupsTaskStep posts the env's wakeups/v1 config as the stored card declares it and records it, and +refuses before any request on a gateway without wakeups/v1. Only the gateway HTTP boundary is mocked.""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from agent_env.env.env import DeployedEnv, DeployedGatewayEnv +from agent_env.env.gateway.constants import EXT_WAKEUPS_URI, GATEWAY_EXTENSIONS, WELL_KNOWN_PATH +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.registry import get_task_step_registry +from agent_env.task_step.task_steps.configure_wakeups import ConfigureWakeupsTaskStep + +_GATEWAY = "http://gw:18765" + + +@pytest.mark.asyncio +async def test_posts_the_config_and_records_what_the_gateway_applied(monkeypatch): + sent = _mock_gateway(monkeypatch, {"scheduled_checks": True, "max_active": 5, "min_interval": "1:00:00"}) + out = await _step(max_active=5).execute(TaskStepContext(deployed_envs=[_record()], metadata={})) + [request] = sent + assert (request.method, str(request.url)) == ("POST", f"{_GATEWAY}/wakeups/config") + assert json.loads(request.content) == {"scheduled_checks": True, "max_active": 5, "min_interval": "PT1H"} + [entry] = out.metadata["wakeup_configurations"] + assert entry["applied"] == {"scheduled_checks": True, "max_active": 5, "min_interval": "1:00:00"} + + +@pytest.mark.asyncio +async def test_a_gateway_without_wakeups_refuses_before_any_request(monkeypatch): + sent = _mock_gateway(monkeypatch, {}) + older = [e for e in GATEWAY_EXTENSIONS if e["uri"] != EXT_WAKEUPS_URI] + with pytest.raises(RuntimeError, match="without urn:agentenv:wakeups/v1"): + await _step().execute(TaskStepContext(deployed_envs=[_record(older)], metadata={})) + assert sent == [] + + +@pytest.mark.asyncio +async def test_a_refused_config_raises_with_the_gateway_s_reason(monkeypatch): + _mock_gateway(monkeypatch, {"ok": False, "error": "['schedule_check'] already registered by a backing server"}, + status=409) + with pytest.raises(RuntimeError, match="HTTP 409.*schedule_check"): + await _step().execute(TaskStepContext(deployed_envs=[_record()], metadata={})) + + +def test_round_trips_and_is_registered(): + step = _step(max_active=3, min_interval="PT2H") + again = ConfigureWakeupsTaskStep.from_dict(step.to_dict()) + assert (again.max_active, again.min_interval, again.scheduled_checks) == (3, "PT2H", True) + assert get_task_step_registry()["configure_wakeups"] is ConfigureWakeupsTaskStep + + +@pytest.mark.parametrize("kwargs", [{"max_active": 0}, {"scheduled_checks": "yes"}]) +def test_bad_settings_are_refused_when_the_step_is_built(kwargs): + with pytest.raises(ValueError): + _step(**kwargs) + + +def _step(**kwargs) -> ConfigureWakeupsTaskStep: + return ConfigureWakeupsTaskStep(id="wakeups-1", version=None, env_id="env-x", **kwargs) + + +def _record(extensions: list = GATEWAY_EXTENSIONS) -> DeployedEnv: + return DeployedGatewayEnv(env_id="env-x", env_version=1, gateway_url=_GATEWAY, mcp_url="", db_web_url=None, + sandbox_id="sb", environment_card_url=f"{_GATEWAY}{WELL_KNOWN_PATH}", + environment_card={"name": "gw", "capabilities": {"extensions": extensions}}) + + +def _mock_gateway(monkeypatch, body: dict, status: int = 200) -> list[httpx.Request]: + sent, real = [], httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(status, json=body) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + return sent diff --git a/tst/unit/task_step/test_register_env_triggers.py b/tst/unit/task_step/test_register_env_triggers.py index 526f47bb2..7c7c533b7 100644 --- a/tst/unit/task_step/test_register_env_triggers.py +++ b/tst/unit/task_step/test_register_env_triggers.py @@ -578,3 +578,18 @@ def handle(request: httpx.Request) -> httpx.Response: monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) return sent + + +@pytest.mark.asyncio +async def test_a_gateway_without_notify_agent_refuses_it(monkeypatch): + sent = _mock_gateway(monkeypatch, {"ok": True, "added": ["n"], "all": ["n"]}) + notify = [{"id": "n", "when": {"type": "action", "tool": "t"}, + "actions": [{"type": "notify_agent", "channel": "Gmail", "text": "New email"}]}] + context = TaskStepContext(deployed_envs=[_record(_with_features("named_executors", "provoking_roles"))], + deployed_agents=[], metadata={}) + with pytest.raises(RuntimeError, match=r"does not declare \['notify_agent'\]"): + await _step(triggers=notify).execute(context) + assert sent == [] + context.deployed_envs = [_record()] + await _step(triggers=notify).execute(context) + assert json.loads(sent[0].content)["triggers"] == notify From 8f9f624b577a534ba7adfb7967eee3461bfa0731 Mon Sep 17 00:00:00 2001 From: Sarah Chitty Date: Tue, 27 Oct 2026 10:07:37 -0700 Subject: [PATCH 2/4] =?UTF-8?q?fix(gateway):=20wakeups=20=E2=80=94=20count?= =?UTF-8?q?=20repeats=20instead=20of=20stepping,=20judge=20due=20on=20time?= =?UTF-8?q?s,=20wait=20for=20discovery?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - take counts how often a repeating check came due and its next time directly, so a long jump over a short every costs nothing, and works the whole batch out before marking anything handed out - schedule refuses an every (or at) too long to count forward, which would otherwise break every later take - notifications are due by their parsed time, not their text (fractional seconds sorted wrong) - /wakeups/config finishes tool discovery before its name check (503 if discovery fails), so a later discovery can't replace schedule_check unseen - configure_wakeups logs a one-line completion, not the request body Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/gateway.py | 5 ++ src/agent_env/env/gateway/wakeups.py | 80 ++++++++++++------- .../task_step/task_steps/configure_wakeups.py | 2 +- tst/unit/env/gateway/wakeups_test.py | 43 ++++++++++ 4 files changed, 102 insertions(+), 28 deletions(-) diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 219e6541f..c8c4a517e 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -312,6 +312,11 @@ async def _wakeup_write(request: Request, method) -> Response: @self._mcp.custom_route("/wakeups/config", methods=["POST"]) async def wakeups_config(request: Request) -> Response: + try: # the backing servers' tools must be known, or a later discovery would replace ours unseen + await self._ensure_tools_discovered() + except Exception as e: + return self._json_response( + {"ok": False, "error": f"tool discovery has not completed ({type(e).__name__}); retry"}, 503) taken = sorted(set(WAKEUP_TOOL_NAMES) & set(self._tool_server_urls)) if taken: return self._json_response( diff --git a/src/agent_env/env/gateway/wakeups.py b/src/agent_env/env/gateway/wakeups.py index 7872c137b..eacb36167 100644 --- a/src/agent_env/env/gateway/wakeups.py +++ b/src/agent_env/env/gateway/wakeups.py @@ -70,6 +70,8 @@ def _when(text: Any, now: datetime, field: str) -> datetime: return now + _parse_duration(text) if text.startswith("P") else _parse_rfc3339(text) except ClockError as e: raise WakeupError(f"{field}: {e}") + except OverflowError: + raise WakeupError(f"{field} is too far away") class Wakeups: @@ -102,7 +104,7 @@ def configure(self, body: Any) -> dict: raise WakeupError("max_active must be a positive integer") try: min_interval = _parse_duration(body.get("min_interval", DEFAULT_MIN_INTERVAL)) - except ClockError as e: + except (ClockError, OverflowError) as e: raise WakeupError(f"min_interval: {e}") if min_interval <= timedelta(0): raise WakeupError("min_interval must be a positive duration") @@ -145,13 +147,18 @@ def schedule(self, note: Any, at: Any = None, every: Any = None, until: Any = No if every is not None: try: step = _parse_duration(every) - except ClockError as e: + except (ClockError, OverflowError) as e: raise WakeupError(f"every: {e}") if step < self._min_interval: raise WakeupError(f"every must be at least {self._min_interval} apart") - first = _when(at, now, "at") if at is not None else (now + step if step else None) - if first is None: + if at is None and step is None: raise WakeupError("give at (when) or every (how often)") + first = _when(at, now, "at") if at is not None else _when(every, now, "every") + if step is not None: + try: + first + step # a repeat that can't be counted forward would break every later take + except OverflowError: + raise WakeupError("every is too long") if first <= now: raise WakeupError(f"at {_iso(first)} has already passed; it is now {_iso(now)}") end = _when(until, now, "until") if until is not None else None @@ -195,16 +202,16 @@ def notify(self, channel: str, text: str, trigger_id: str | None = None) -> dict now = self._clock.now() with self._lock: item = {"id": self._next_id("note"), "channel": channel, "text": text, "trigger_id": trigger_id, - "at": _iso(now) if now else None, "status": "pending"} + "at": _iso(now) if now else None, "status": "pending", "_at": now} self._notifications.append(item) - return dict(item) + return _public_note(item) # ---- harness ---- def state(self) -> dict: with self._lock: active = [c for c in self._checks.values() if c["status"] == "active"] - pending = [dict(n) for n in self._notifications if n["status"] == "pending"] + pending = [_public_note(n) for n in self._notifications if n["status"] == "pending"] nxt = min((c["next"] for c in active), default=None) return {"scheduled_checks": self._enabled, "max_active": self._max_active, "min_interval": str(self._min_interval), "next_check_at": _iso(nxt) if nxt else None, @@ -218,28 +225,47 @@ def take(self, body: Any = None) -> dict: body = body if isinstance(body, dict) else {} now = self._clock.now() until = _when(body["until"], now, "until") if body.get("until") is not None and now else now - items: list[dict] = [] + due: list[tuple[datetime | None, dict]] = [] with self._lock: - for check in self._checks.values(): - if check["status"] != "active" or until is None or check["next"] > until: - continue - due, times = check["next"], 0 - while check["status"] == "active" and check["next"] <= until: - times += 1 - if check["every"] and (check["until"] is None or check["next"] + check["every"] <= check["until"]): - check["next"] = check["next"] + check["every"] - else: - check["status"] = "done" + # work the whole batch out before marking anything handed out + checks = [(c, *_after(c, until)) for c in self._checks.values() + if c["status"] == "active" and until is not None and c["next"] <= until] + notes = [n for n in self._notifications + if n["status"] == "pending" and (until is None or n["_at"] is None or n["_at"] <= until)] + for check, times, nxt in checks: + due.append((check["next"], {"kind": "scheduled_check", "check_id": check["check_id"], + "note": check["note"], "at": _iso(check["next"]), "times": times})) check["fired"] += 1 - items.append({"kind": "scheduled_check", "check_id": check["check_id"], "note": check["note"], - "at": _iso(due), "times": times}) - for note in self._notifications: - if note["status"] == "pending" and (until is None or note["at"] is None or note["at"] <= _iso(until)): - note["status"] = "delivered" - items.append({"kind": "notification", "id": note["id"], "channel": note["channel"], - "text": note["text"], "at": note["at"], "trigger_id": note["trigger_id"]}) - items.sort(key=lambda i: i["at"] or "") - return {"items": items, "until": _iso(until) if until else None} + if nxt is None: + check["status"] = "done" + else: + check["next"] = nxt + for note in notes: + note["status"] = "delivered" + due.append((note["_at"], {"kind": "notification", "id": note["id"], "channel": note["channel"], + "text": note["text"], "at": note["at"], "trigger_id": note["trigger_id"]})) + due.sort(key=lambda d: (d[0] is not None, d[0])) + return {"items": [item for _, item in due], "until": _iso(until) if until else None} + + +def _after(check: dict, until: datetime) -> tuple[int, datetime | None]: + """How many times a due check came due by `until`, and its next time after that (None: it is done). Counted, + not stepped, so a long jump over a short `every` costs nothing.""" + step, end = check["every"], check["until"] + if not step: + return 1, None + times = (until - check["next"]) // step + 1 + if end is not None: + times = min(times, (end - check["next"]) // step + 1) + try: + nxt = check["next"] + step * times + except OverflowError: + return times, None + return times, nxt if end is None or nxt <= end else None + + +def _public_note(note: dict) -> dict: + return {k: v for k, v in note.items() if not k.startswith("_")} def _tool_error(text: str) -> CallToolResult: diff --git a/src/agent_env/task_step/task_steps/configure_wakeups.py b/src/agent_env/task_step/task_steps/configure_wakeups.py index 0700c37f2..b66a2506f 100644 --- a/src/agent_env/task_step/task_steps/configure_wakeups.py +++ b/src/agent_env/task_step/task_steps/configure_wakeups.py @@ -82,7 +82,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: state = await deployed.invoke(EXT_WAKEUPS_URI, "config", body, timeout=self.timeout_seconds) except httpx.HTTPStatusError as e: raise RuntimeError(f"wakeups config failed (HTTP {e.response.status_code}): {e.response.text}") from e - logger.info(f"configure_wakeups env={self.env_id}: {body}") + logger.info(f"configure_wakeups completed for env={self.env_id}") context.metadata.setdefault("wakeup_configurations", []).append({ "step_id": self.id, "env_id": self.env_id, **body, "applied": {k: state.get(k) for k in ("scheduled_checks", "max_active", "min_interval")}, diff --git a/tst/unit/env/gateway/wakeups_test.py b/tst/unit/env/gateway/wakeups_test.py index 85f2c35b8..2d9134bd2 100644 --- a/tst/unit/env/gateway/wakeups_test.py +++ b/tst/unit/env/gateway/wakeups_test.py @@ -66,6 +66,8 @@ def test_a_check_takes_a_duration_from_now_or_a_time_with_its_offset(): ({"note": "x", "every": "PT10M"}, "at least"), ({"note": "x", "at": "PT1H", "until": "P1D"}, "until only applies"), ({"note": "x", "every": "PT6H", "until": "PT1H"}, "before the first"), + ({"note": "x", "at": "PT1H", "every": "P3000000D"}, "too long"), + ({"note": "x", "at": "P3000000D"}, "too far away"), ]) def test_bad_checks_are_refused_with_the_reason(kwargs, error): with pytest.raises(WakeupError, match=error): @@ -108,6 +110,16 @@ def test_a_repeating_check_moves_past_a_long_jump_and_says_how_often_it_came_due assert w.list()["checks"][0]["next"] == "2027-02-02T20:00:00Z" +def test_a_jump_over_a_short_every_is_counted_not_stepped(): + clock = _frozen_clock() + w = _wakeups(clock, min_interval="PT0.001S") + w.schedule("tick", every="PT0.001S") + clock.set_time("2027-02-02T15:00:00Z", 0) + [item] = w.take({})["items"] + assert item["times"] == 86_400_000 + assert w.list()["checks"][0]["next"] == "2027-02-02T15:00:00.001000Z" + + def test_a_repeating_check_ends_at_until(): clock = _frozen_clock() w = _wakeups(clock) @@ -135,6 +147,15 @@ def test_notifications_queue_at_the_virtual_time_and_are_taken_once(): assert w.take({})["items"] == [] and w.state()["pending_notifications"] == [] +def test_due_is_judged_on_the_time_not_its_text(): + clock = _frozen_clock() + w = _wakeups(clock) + w.notify("Messages", "on the second") # 15:00:00Z sorts after 15:00:00.500000Z as text + w.schedule("half a second on", at="PT0.5S") + items = w.take({"until": "PT0.75S"})["items"] + assert [i.get("text") or i.get("note") for i in items] == ["on the second", "half a second on"] + + def test_clear_turns_checks_off_and_forgets_everything(): w = _wakeups() w.schedule("x", at="PT1H") @@ -202,6 +223,28 @@ def test_config_is_refused_over_a_backing_server_tool_of_the_same_name(gateway): assert r.status_code == 409 and "schedule_check" in r.json()["error"] +def test_config_waits_for_tool_discovery_and_sees_what_it_found(gateway, monkeypatch): + gateway._tools_discovered = False + + async def discover(): + gateway._tool_server_urls["schedule_check"] = "http://svc:1/mcp" + gateway._tools_discovered = True + + monkeypatch.setattr(gateway, "_ensure_tools_discovered", discover) + r = _request(gateway, "POST", "/wakeups/config", {}) + assert r.status_code == 409 and "schedule_check" in r.json()["error"] + + +def test_config_is_a_503_while_tool_discovery_fails(gateway, monkeypatch): + async def discover(): + raise RuntimeError("backing server down") + + monkeypatch.setattr(gateway, "_ensure_tools_discovered", discover) + r = _request(gateway, "POST", "/wakeups/config", {}) + assert r.status_code == 503 and "retry" in r.json()["error"] + assert not set(WAKEUP_TOOL_NAMES) & set(gateway._mcp._tool_manager._tools) + + def test_a_bad_config_is_a_400(gateway): assert _request(gateway, "POST", "/wakeups/config", {"max_active": -1}).status_code == 400 From 0878659b07cf2728beeb3fdaa25ffd696471432f Mon Sep 17 00:00:00 2001 From: Sarah Chitty Date: Tue, 27 Oct 2026 13:37:42 -0700 Subject: [PATCH 3/4] fix(gateway): wakeups echo durations as ISO 8601 (PT3H), the form the tools and config take The state, list_scheduled_checks and the too-short error showed Python's timedelta text (3:00:00), which /wakeups/config then refused when sent back. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/wakeups.py | 15 ++++++++++++--- tst/unit/env/gateway/wakeups_test.py | 12 ++++++++++++ tst/unit/task_step/test_configure_wakeups.py | 4 ++-- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/src/agent_env/env/gateway/wakeups.py b/src/agent_env/env/gateway/wakeups.py index eacb36167..7f0b5697f 100644 --- a/src/agent_env/env/gateway/wakeups.py +++ b/src/agent_env/env/gateway/wakeups.py @@ -150,7 +150,7 @@ def schedule(self, note: Any, at: Any = None, every: Any = None, until: Any = No except (ClockError, OverflowError) as e: raise WakeupError(f"every: {e}") if step < self._min_interval: - raise WakeupError(f"every must be at least {self._min_interval} apart") + raise WakeupError(f"every must be at least {_iso_duration(self._min_interval)} apart") if at is None and step is None: raise WakeupError("give at (when) or every (how often)") first = _when(at, now, "at") if at is not None else _when(every, now, "every") @@ -190,7 +190,7 @@ def cancel(self, check_id: Any) -> dict: def _public(check: dict) -> dict: out = {"check_id": check["check_id"], "note": check["note"], "next": _iso(check["next"])} if check["every"]: - out["every"] = str(check["every"]) + out["every"] = _iso_duration(check["every"]) if check["until"]: out["until"] = _iso(check["until"]) return out @@ -214,7 +214,7 @@ def state(self) -> dict: pending = [_public_note(n) for n in self._notifications if n["status"] == "pending"] nxt = min((c["next"] for c in active), default=None) return {"scheduled_checks": self._enabled, "max_active": self._max_active, - "min_interval": str(self._min_interval), "next_check_at": _iso(nxt) if nxt else None, + "min_interval": _iso_duration(self._min_interval), "next_check_at": _iso(nxt) if nxt else None, "checks": [{**self._public(c), "status": c["status"], "fired": c["fired"]} for c in self._checks.values()], "pending_notifications": pending} @@ -264,6 +264,15 @@ def _after(check: dict, until: datetime) -> tuple[int, datetime | None]: return times, nxt if end is None or nxt <= end else None +def _iso_duration(td: timedelta) -> str: + """A duration in the ISO 8601 form the tools take (PT3H, P1DT6H, PT0.5S), so what they echo can be sent back.""" + hours, rest = divmod(td.seconds, 3600) + minutes, seconds = divmod(rest, 60) + secs = f"{seconds}.{td.microseconds:06d}".rstrip("0") if td.microseconds else (str(seconds) if seconds else "") + clock = (f"{hours}H" if hours else "") + (f"{minutes}M" if minutes else "") + (f"{secs}S" if secs else "") + return "P" + (f"{td.days}D" if td.days else "") + (f"T{clock}" if clock else "") if td.days or clock else "PT0S" + + def _public_note(note: dict) -> dict: return {k: v for k, v in note.items() if not k.startswith("_")} diff --git a/tst/unit/env/gateway/wakeups_test.py b/tst/unit/env/gateway/wakeups_test.py index 2d9134bd2..020051b64 100644 --- a/tst/unit/env/gateway/wakeups_test.py +++ b/tst/unit/env/gateway/wakeups_test.py @@ -74,6 +74,18 @@ def test_bad_checks_are_refused_with_the_reason(kwargs, error): _wakeups().schedule(**kwargs) +@pytest.mark.parametrize("text", ["PT3H", "P1D", "P1DT6H30M", "PT0.001S", "PT1M0.5S"]) +def test_durations_are_echoed_in_the_form_the_tools_take(text): + w = _wakeups(min_interval="PT0.001S") + assert w.schedule("x", every=text)["every"] == text + assert _wakeups(min_interval=text).state()["min_interval"] == text + + +def test_the_too_short_error_names_the_minimum_as_a_duration_to_send_back(): + with pytest.raises(WakeupError, match=r"at least PT3H apart"): + _wakeups(min_interval="PT3H").schedule("x", every="PT1H") + + def test_the_active_cap_refuses_one_more_until_one_is_cancelled(): w = _wakeups(max_active=2) first = w.schedule("a", at="PT1H")["check_id"] diff --git a/tst/unit/task_step/test_configure_wakeups.py b/tst/unit/task_step/test_configure_wakeups.py index ff04738f1..1ab1941a4 100644 --- a/tst/unit/task_step/test_configure_wakeups.py +++ b/tst/unit/task_step/test_configure_wakeups.py @@ -19,13 +19,13 @@ @pytest.mark.asyncio async def test_posts_the_config_and_records_what_the_gateway_applied(monkeypatch): - sent = _mock_gateway(monkeypatch, {"scheduled_checks": True, "max_active": 5, "min_interval": "1:00:00"}) + sent = _mock_gateway(monkeypatch, {"scheduled_checks": True, "max_active": 5, "min_interval": "PT1H"}) out = await _step(max_active=5).execute(TaskStepContext(deployed_envs=[_record()], metadata={})) [request] = sent assert (request.method, str(request.url)) == ("POST", f"{_GATEWAY}/wakeups/config") assert json.loads(request.content) == {"scheduled_checks": True, "max_active": 5, "min_interval": "PT1H"} [entry] = out.metadata["wakeup_configurations"] - assert entry["applied"] == {"scheduled_checks": True, "max_active": 5, "min_interval": "1:00:00"} + assert entry["applied"] == {"scheduled_checks": True, "max_active": 5, "min_interval": "PT1H"} @pytest.mark.asyncio From 0965e7c6f4d027d2623c3bf0686398a390fd9de4 Mon Sep 17 00:00:00 2001 From: Sarah Chitty Date: Tue, 27 Oct 2026 15:02:45 -0700 Subject: [PATCH 4/4] =?UTF-8?q?fix(gateway):=20wakeups=20=E2=80=94=20turni?= =?UTF-8?q?ng=20checks=20off=20cancels=20the=20pending=20ones;=20history?= =?UTF-8?q?=20is=20bounded?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - configure with scheduled_checks false cancels the active checks, since the agent no longer has the tool to - delivered notifications are dropped, and state keeps only the last 50 finished or cancelled checks Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/wakeups.py | 15 +++++++++++++++ tst/unit/env/gateway/wakeups_test.py | 22 ++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/src/agent_env/env/gateway/wakeups.py b/src/agent_env/env/gateway/wakeups.py index 7f0b5697f..7bc045499 100644 --- a/src/agent_env/env/gateway/wakeups.py +++ b/src/agent_env/env/gateway/wakeups.py @@ -29,6 +29,7 @@ DEFAULT_MIN_INTERVAL = "PT1H" _MAX_NOTE_CHARS = 2000 _MAX_CHANNEL_CHARS = 60 +_HISTORY = 50 # finished and cancelled checks state() still shows SCHEDULE_CHECK_DESCRIPTION = ( "Schedule a check-in for yourself. You only run when you are messaged, so to look at something again later " @@ -110,6 +111,11 @@ def configure(self, body: Any) -> dict: raise WakeupError("min_interval must be a positive duration") with self._lock: self._enabled, self._max_active, self._min_interval = enabled, max_active, min_interval + if not enabled: # the agent loses its tools, so the checks it can no longer cancel go too + for check in self._checks.values(): + if check["status"] == "active": + check["status"] = "cancelled" + self._prune() return self.state() def clear(self) -> dict: @@ -184,6 +190,7 @@ def cancel(self, check_id: Any) -> dict: if check is None or check["status"] != "active": raise WakeupError(f"no active scheduled check {check_id!r}") check["status"] = "cancelled" + self._prune() return {"cancelled": check_id} @staticmethod @@ -219,6 +226,13 @@ def state(self) -> dict: for c in self._checks.values()], "pending_notifications": pending} + def _prune(self) -> None: + """Drop delivered notifications and all but the last _HISTORY finished checks; the caller holds the lock.""" + finished = [cid for cid, c in self._checks.items() if c["status"] != "active"] + for cid in finished[:max(0, len(finished) - _HISTORY)]: + del self._checks[cid] + self._notifications = [n for n in self._notifications if n["status"] == "pending"] + def take(self, body: Any = None) -> dict: """Everything due by `until` (default: now), oldest first; each is handed out once. A repeating check moves to its next time after `until` and says how many times it came due, so a long jump wakes the agent once.""" @@ -244,6 +258,7 @@ def take(self, body: Any = None) -> dict: note["status"] = "delivered" due.append((note["_at"], {"kind": "notification", "id": note["id"], "channel": note["channel"], "text": note["text"], "at": note["at"], "trigger_id": note["trigger_id"]})) + self._prune() due.sort(key=lambda d: (d[0] is not None, d[0])) return {"items": [item for _, item in due], "until": _iso(until) if until else None} diff --git a/tst/unit/env/gateway/wakeups_test.py b/tst/unit/env/gateway/wakeups_test.py index 020051b64..be0a13db7 100644 --- a/tst/unit/env/gateway/wakeups_test.py +++ b/tst/unit/env/gateway/wakeups_test.py @@ -168,6 +168,28 @@ def test_due_is_judged_on_the_time_not_its_text(): assert [i.get("text") or i.get("note") for i in items] == ["on the second", "half a second on"] +def test_turning_checks_off_cancels_the_ones_still_to_come(): + clock = _frozen_clock() + w = _wakeups(clock) + w.schedule("sweep", every="P1D") + w.configure({"scheduled_checks": False}) + clock.set_time("2027-02-05T15:00:00Z", 0) + assert w.state()["next_check_at"] is None and w.take({})["items"] == [] + assert [c["status"] for c in w.state()["checks"]] == ["cancelled"] + + +def test_history_is_bounded_and_delivered_notifications_are_dropped(): + from agent_env.env.gateway import wakeups as wakeups_mod + w = _wakeups() + for n in range(wakeups_mod._HISTORY + 10): + w.cancel(w.schedule(f"check {n}", at="PT1H")["check_id"]) + checks = w.state()["checks"] + assert len(checks) == wakeups_mod._HISTORY and checks[-1]["note"] == f"check {wakeups_mod._HISTORY + 9}" + w.notify("Messages", "hi") + w.take({}) + assert w._notifications == [] + + def test_clear_turns_checks_off_and_forgets_everything(): w = _wakeups() w.schedule("x", at="PT1H")