diff --git a/.github/scripts/check_plugin_api.py b/.github/scripts/check_plugin_api.py index b0ee5462..99a86066 100644 --- a/.github/scripts/check_plugin_api.py +++ b/.github/scripts/check_plugin_api.py @@ -34,7 +34,7 @@ "agent_env.explorer.plugin.ExplorerPlugin", ) USED = ( - "agent_env.task_step.context.TaskStepContext", "agent_env.plugins", + "agent_env.task_step.context.TaskStepContext", "agent_env.task_step.context.HEARTBEAT_KEY", "agent_env.plugins", "agent_env.providers.env_providers.env_provider.build_env_provider", "agent_env.env.store.register_env_instance", # What a plugin reads of recorded runs. "agent_env.task.store.task_instances", "agent_env.task.store.count_task_instances", diff --git a/AGENTS.md b/AGENTS.md index bca3cd73..a3681398 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -180,6 +180,9 @@ orchestrators call it the same way. What callers rely on: a context with `TaskStepContext.to_dict()`, not `dataclasses.asdict`. - A run is resumable: pass a previously persisted `context`, a `start_step` and the existing `instance_id`; completed steps are seeded as SUCCESS. +- A runner that must hear from a long step puts a zero-argument callable at + `context.metadata[HEARTBEAT_KEY]` (`task_step/context.py`); a step that blocks for long calls + `context.heartbeat()`, a no-op without one. The key is in `_REDACTED_KEYS`, so it is never persisted. - An episode (`EpisodeSpec`, `task/episode/`) is one `Task.run()` of a copy of the task in which the `deploy_agent` and `prompt_agent` steps of each of its `trainer_agents` are the plain-data `episode_agent` and `episode_turn`, which wait for the caller instead of deploying or prompting. Its diff --git a/src/agent_env/task_step/context.py b/src/agent_env/task_step/context.py index 7335d461..936d00dd 100644 --- a/src/agent_env/task_step/context.py +++ b/src/agent_env/task_step/context.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextlib import dataclasses from dataclasses import dataclass, field from typing import Any @@ -10,8 +11,11 @@ from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, host_url_for from agent_env.providers.sandbox_providers.sandbox_provider import reachable_url +HEARTBEAT_KEY = "_heartbeat_fn" + _REDACTED_KEYS = { "litellm_api_key", "judge_litellm_api_key", "usersim_api_key", "remote_tokens", "cf_access_client_secret", + HEARTBEAT_KEY, } @@ -222,6 +226,14 @@ def to_dict(self) -> dict[str, Any]: d["prompt_responses"] = [response.to_dict() for response in self.prompt_responses] return d + def heartbeat(self) -> None: + """Tell the runner this step is still alive, through the callable it put at ``metadata[HEARTBEAT_KEY]``. + A no-op without one; a beat that raises is dropped, never failing the step.""" + beat = self.metadata.get(HEARTBEAT_KEY) + if callable(beat): + with contextlib.suppress(Exception): + beat() + def record_model(self, role: str, step_id: str, model: str | None) -> None: """Record the model ``step_id`` ran ``role`` on, None when the agent chose its own, for the run's summary.""" self.metadata.setdefault("models", {}).setdefault(role, {})[step_id] = model diff --git a/src/agent_env/task_step/task_steps/review.py b/src/agent_env/task_step/task_steps/review.py index 2df069d3..ac270566 100644 --- a/src/agent_env/task_step/task_steps/review.py +++ b/src/agent_env/task_step/task_steps/review.py @@ -1,8 +1,8 @@ """Human-in-the-loop review: pause the task until an operator continues or aborts. -Registers an ``awaiting`` review and blocks (heartbeating to keep the activity -alive) until the hub records a decision or the timeout elapses — continue returns -normally, abort raises so downstream steps don't run. The wait holds no sandbox, +Registers an ``awaiting`` review and blocks (beating ``context.heartbeat()`` so the +runner keeps the step alive) until the hub records a decision or the timeout elapses — +continue returns normally, abort raises so downstream steps don't run. The wait holds no sandbox, so place it after the agent's output is persisted and before ``collect_artifacts`` (an abort then excludes the run's artifacts). """ @@ -113,12 +113,11 @@ async def _await_decision(self, store, context: TaskStepContext) -> dict: """Poll until decided (return the record) or the wall-clock deadline (raise TimeoutError).""" from agent_env.task_step.review_store import ABORT, CONTINUE - heartbeat = context.metadata.get("_heartbeat_fn") loop = asyncio.get_running_loop() deadline = loop.time() + self.timeout_seconds # wall-clock, so a non-positive interval can't spin forever interval = self.poll_interval_seconds if self.poll_interval_seconds > 0 else DEFAULT_POLL_INTERVAL_SECONDS while loop.time() < deadline: - self._beat(heartbeat) + context.heartbeat() record = store.get(context.instance_id, self.id) or {} if record.get("state") in (CONTINUE, ABORT): return record @@ -142,11 +141,3 @@ def _apply_decision(self, context: TaskStepContext, record: dict) -> TaskStepCon raise ReviewAbortedError( f"review '{self.id}' aborted by operator" + (f" ({decided_by})" if decided_by else "") ) - - @staticmethod - def _beat(heartbeat) -> None: - if callable(heartbeat): - try: - heartbeat() - except Exception: # best-effort; never fail the wait on it - pass diff --git a/tst/unit/task_step/context_ops_journal_test.py b/tst/unit/task_step/context_ops_journal_test.py index a84475f7..3fad38f7 100644 --- a/tst/unit/task_step/context_ops_journal_test.py +++ b/tst/unit/task_step/context_ops_journal_test.py @@ -5,7 +5,7 @@ from __future__ import annotations from agent_env.task.task import _rebuild_context -from agent_env.task_step.context import TaskStepContext, regraft_redacted_keys +from agent_env.task_step.context import HEARTBEAT_KEY, TaskStepContext, regraft_redacted_keys from agent_env.task_step.context_ops import ContextUpdateOps, build_context_update_ops @@ -83,3 +83,15 @@ def test_the_judge_key_never_reaches_the_store_but_survives_a_rebuild(): assert "sk-judge" not in repr(build_context_update_ops(None, live).to_journal_dict()) _rebuild_context(live, stored) assert live.metadata["user_overrides"] == {"judge_litellm_api_key": "sk-judge", "agent_effort": "high"} + + +def test_the_runner_heartbeat_is_never_persisted_but_survives_a_rebuild(): + def beat(): + return None + + live = TaskStepContext(metadata={HEARTBEAT_KEY: beat, "workflow_id": "w"}) + stored = live.to_safe_dict() + assert stored["metadata"] == {"workflow_id": "w"} + assert HEARTBEAT_KEY not in repr(build_context_update_ops(None, live).to_journal_dict()) + _rebuild_context(live, stored) + assert live.metadata[HEARTBEAT_KEY] is beat diff --git a/tst/unit/task_step/test_review.py b/tst/unit/task_step/test_review.py index 4674f7f0..229a2c4f 100644 --- a/tst/unit/task_step/test_review.py +++ b/tst/unit/task_step/test_review.py @@ -10,6 +10,7 @@ from agent_env.task_step.review_store import ABORT, AWAITING, CONTINUE, ReviewStore from agent_env.config import set_document_store +from agent_env.task_step.context import HEARTBEAT_KEY, TaskStepContext from agent_env.task_step.task_steps.review import ( ReviewAbortedError, ReviewTaskStep, @@ -33,10 +34,10 @@ def get(self, instance_id, step_id): return {"state": state, "decided_by": "op@example.com"} -def _ctx(prompt_responses=None): - return SimpleNamespace( +def _ctx(prompt_responses=None, heartbeat=lambda: None): + return TaskStepContext( instance_id="inst-1", - metadata={"_heartbeat_fn": lambda: None}, + metadata={HEARTBEAT_KEY: heartbeat}, prompt_responses=prompt_responses or [], ) @@ -75,6 +76,23 @@ async def test_abort_raises_and_never_returns_context(): await _step().execute(_ctx()) +@pytest.mark.asyncio +async def test_the_wait_beats_the_runner_heartbeat_and_survives_one_that_raises(): + beats = [] + + def beat(): + beats.append(1) + raise ConnectionError("runner unreachable") + + store = FakeStore([AWAITING, AWAITING, CONTINUE]) + with patch( + "agent_env.task_step.review_store.get_review_store", return_value=store + ): + await _step().execute(_ctx(heartbeat=beat)) + + assert len(beats) == 3 + + @pytest.mark.asyncio async def test_timeout_raises(): store = FakeStore([AWAITING])