From 631d36a9ee58b5dfffa04015c7917db63a027b37 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 21 Oct 2026 22:53:04 -0700 Subject: [PATCH 1/5] feat(task-step): derive snapshot_agent_state's artifact id per run when omitted snapshot_agent_state took a fixed artifact_id from the task definition, so every run of one task wrote a new version of the same artifact and a reader of the latest version could get another run's agent state. artifact_id is now optional. When it is omitted the step snapshots into __, the same instance-scoped base SnapshotSeries derives its ids from: distinct across runs, stable across retries of one instance. An explicit artifact_id behaves as before. from_dict tolerates a missing key, so a stored step without one loads; older releases cannot load such a step. FD-4006 Co-Authored-By: Claude Opus 5.5 --- .../snapshot_utils/snapshot_series.py | 4 +- .../task_steps/snapshot_agent_state.py | 20 ++++- .../task_step/test_snapshot_agent_state.py | 78 ++++++++++++++++++- 3 files changed, 94 insertions(+), 8 deletions(-) diff --git a/src/agent_env/task_step/snapshot_utils/snapshot_series.py b/src/agent_env/task_step/snapshot_utils/snapshot_series.py index fe9628d2..f0b8817e 100644 --- a/src/agent_env/task_step/snapshot_utils/snapshot_series.py +++ b/src/agent_env/task_step/snapshot_utils/snapshot_series.py @@ -85,7 +85,7 @@ def from_dict(cls, data: dict) -> "SnapshotConfig": ) -def _rollout_base(step_id: str, instance_id: Optional[str]) -> str: +def rollout_base(step_id: str, instance_id: Optional[str]) -> str: """The id every artifact id in one series is derived from: the rollout's instance id. Instance-scoped, so concurrent rollouts never contend on one id and retries @@ -132,7 +132,7 @@ def __init__( self.a2a_context_id = a2a_context_id self.config = config self.trajectory_output_prefix = trajectory_output_prefix - self._base = _rollout_base(step_id, instance_id) + self._base = rollout_base(step_id, instance_id) self.workspace_artifact_id = derive_id(self._base, f"{step_id}-workspace") if is_local_id(self._base): validate_local_id(derive_id(self._base, f"snapshot-{step_id}-workspace")) diff --git a/src/agent_env/task_step/task_steps/snapshot_agent_state.py b/src/agent_env/task_step/task_steps/snapshot_agent_state.py index 459af7b3..7b52cf77 100644 --- a/src/agent_env/task_step/task_steps/snapshot_agent_state.py +++ b/src/agent_env/task_step/task_steps/snapshot_agent_state.py @@ -6,10 +6,12 @@ import logging from typing import ClassVar, Optional +from agent_env.store.ids import derive_id, is_local_id, validate_local_id from agent_env.task_step.context import TaskStepContext from agent_env.entity.refs import EntityRef, RefRole from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.task_step.snapshot_utils import agent_state_capture +from agent_env.task_step.snapshot_utils.snapshot_series import rollout_base logger = logging.getLogger(__name__) @@ -26,6 +28,10 @@ class SnapshotAgentStateTaskStep(TaskStep): ``{id, version, bundle_object_url, source_agent_name, source_context_id}`` to ``context.metadata['agent_snapshots']`` (a list of such entries). + Without an ``artifact_id`` the artifact is per run, derived from the instance id + (``__``), so concurrent runs of one task never version each + other's snapshot and a retry of one instance reuses its id. + When ``env_id`` is set, the step ALSO reads each MCP service's current state from the deployed env (its ``data/get`` JSON, else its ``/export-state``), uploads each as JSON alongside the workspace tarball, and @@ -49,7 +55,7 @@ def __init__( self, id: str, version: Optional[int], - artifact_id: str, + artifact_id: Optional[str], prompt_id: str, agent_name: Optional[str] = None, env_id: Optional[str] = None, @@ -84,7 +90,7 @@ def to_dict(self) -> dict: def from_dict(cls, data: dict) -> SnapshotAgentStateTaskStep: return cls( **cls._base_from_dict(data), - artifact_id=data["artifact_id"], + artifact_id=data.get("artifact_id"), prompt_id=data["prompt_id"], agent_name=data.get("agent_name"), env_id=data.get("env_id"), @@ -92,6 +98,14 @@ def from_dict(cls, data: dict) -> SnapshotAgentStateTaskStep: timeout_seconds=data.get("timeout_seconds", cls.DEFAULT_TIMEOUT_SECONDS), ) + def resolve_artifact_id(self, context: TaskStepContext) -> str: + if self.artifact_id: + return self.artifact_id + artifact_id = derive_id(rollout_base(self.id, context.instance_id), self.id) + if is_local_id(artifact_id): + validate_local_id(artifact_id) + return artifact_id + async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if deployed is None: @@ -115,7 +129,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: a2a_card=deployed.a2a_card or {}, agent_name=self.agent_name, a2a_context_id=a2a_context_id, - artifact_id=self.artifact_id, + artifact_id=self.resolve_artifact_id(context), timeout_seconds=self.timeout_seconds, sandbox_type=deployed.sandbox_type, ) diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index 08a46d70..22cb2644 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -1,5 +1,6 @@ -"""``_capture_universe_state`` publishes unsigned object URLs to -``context.metadata['snapshot_json_url']``, and reads each child env's route from the stored env card.""" +"""``snapshot_agent_state``: the artifact it snapshots into, and ``_capture_universe_state``, which publishes +unsigned object URLs to ``context.metadata['snapshot_json_url']`` and reads each child env's route from the +stored env card.""" from __future__ import annotations import json @@ -12,7 +13,8 @@ from agent_env.env.env import DeployedGatewayEnv from agent_env.store import set_object_store from agent_env.store.object_store import DEFAULT_CONTENT_TYPE -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import DeployedAgent, PromptResponse, TaskStepContext +from agent_env.task_step.snapshot_utils import agent_state_capture from agent_env.task_step.task_steps import snapshot_agent_state as mod from agentenv_protocol import client as protocol_v1 from tst.unit.event_loop_probe import on_event_loop @@ -261,3 +263,73 @@ async def _state(_deployed, _gateway, name): name: local.object_url(f"agent_snapshots/oc/3-deadbeef/services/{name}.json") for name in ("calendar", "contacts") } assert json.loads(local.get(published["calendar"])) == {"service": "calendar"} + + +def _per_run_step(artifact_id=None): + return mod.SnapshotAgentStateTaskStep( + id="t-agent-snapshot", version=None, artifact_id=artifact_id, prompt_id="main", agent_name="agent", + ) + + +def _run_ctx(instance_id): + ctx = TaskStepContext(instance_id=instance_id) + ctx.deployed_agents = [DeployedAgent(agent_name="agent", api_url="https://agent")] + ctx.prompt_responses = [PromptResponse(prompt_id="main", response="ok", a2a_context_id="ctx-1")] + return ctx + + +@pytest.fixture +def captured_into(monkeypatch): + artifact_ids: list[str] = [] + + async def _capture(*, artifact_id, **_): + artifact_ids.append(artifact_id) + return agent_state_capture.WorkspaceCapture( + universe_id=artifact_id, universe_version=1, + bundle_object_url=f"fake://b/{artifact_id}/workspace.tar.gz", capture_prefix=f"fake://b/{artifact_id}/", + ) + + monkeypatch.setattr(agent_state_capture, "capture_workspace", _capture) + return artifact_ids + + +@pytest.mark.asyncio +async def test_without_an_artifact_id_each_run_snapshots_into_its_own_artifact(captured_into): + step = _per_run_step() + first, retry, other = _run_ctx("inst-a"), _run_ctx("inst-a"), _run_ctx("inst-b") + for ctx in (first, retry, other): + await step.execute(ctx) + + assert captured_into == ["inst-a__t-agent-snapshot", "inst-a__t-agent-snapshot", "inst-b__t-agent-snapshot"] + assert first.metadata["agent_snapshots"][0]["id"] == "inst-a__t-agent-snapshot" + assert other.metadata["agent_snapshots"][0]["id"] == "inst-b__t-agent-snapshot" + assert step.artifact_id is None + + +@pytest.mark.asyncio +async def test_an_explicit_artifact_id_is_used_for_every_run(captured_into): + step = _per_run_step("shared-snapshot") + for instance_id in ("inst-a", "inst-b"): + await step.execute(_run_ctx(instance_id)) + + assert captured_into == ["shared-snapshot", "shared-snapshot"] + + +def test_a_local_instance_derives_a_local_artifact_id(): + assert _per_run_step().resolve_artifact_id(_run_ctx("@local/runs/r1")) == "@local/runs/r1__t-agent-snapshot" + + +@pytest.mark.parametrize("artifact_id", [None, "shared-snapshot"]) +def test_round_trips_with_and_without_an_artifact_id(artifact_id): + data = _per_run_step(artifact_id).to_dict() + restored = mod.SnapshotAgentStateTaskStep.from_dict(data) + + assert restored.artifact_id == artifact_id + assert restored.to_dict() == data + + +def test_a_stored_step_without_the_artifact_id_key_loads(): + data = _per_run_step().to_dict() + del data["artifact_id"] + + assert mod.SnapshotAgentStateTaskStep.from_dict(data).artifact_id is None From 5faf2a3e9a3a549fc892a2a62bfdec40808be5a3 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 21 Oct 2026 23:03:37 -0700 Subject: [PATCH 2/5] feat(task-step): record the capturing step on snapshot_agent_state's entry agent_snapshots has several writers (prompt_agent's series rows and any number of snapshot_agent_state steps), so a reader that wants one step's snapshot needs the step id on the entry rather than re-deriving the artifact id. The sibling capture records already carry capture_step_id. Also tests the derived id without an instance id, and keeps the resolver private. FD-4006 Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/snapshot_agent_state.py | 7 ++++--- tst/unit/task_step/test_snapshot_agent_state.py | 8 +++++++- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/agent_env/task_step/task_steps/snapshot_agent_state.py b/src/agent_env/task_step/task_steps/snapshot_agent_state.py index 7b52cf77..eb4f91fa 100644 --- a/src/agent_env/task_step/task_steps/snapshot_agent_state.py +++ b/src/agent_env/task_step/task_steps/snapshot_agent_state.py @@ -25,7 +25,7 @@ class SnapshotAgentStateTaskStep(TaskStep): agent's ``/ext/snapshot`` extension to write the conversation transcript to the object store, then wraps that prefix as a ``FileArtifactUniverse`` via ``put_existing`` (no download/re-upload). Appends an entry recording - ``{id, version, bundle_object_url, source_agent_name, source_context_id}`` to + ``{id, version, bundle_object_url, source_agent_name, source_context_id, capture_step_id}`` to ``context.metadata['agent_snapshots']`` (a list of such entries). Without an ``artifact_id`` the artifact is per run, derived from the instance id @@ -98,7 +98,7 @@ def from_dict(cls, data: dict) -> SnapshotAgentStateTaskStep: timeout_seconds=data.get("timeout_seconds", cls.DEFAULT_TIMEOUT_SECONDS), ) - def resolve_artifact_id(self, context: TaskStepContext) -> str: + def _resolve_artifact_id(self, context: TaskStepContext) -> str: if self.artifact_id: return self.artifact_id artifact_id = derive_id(rollout_base(self.id, context.instance_id), self.id) @@ -129,7 +129,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: a2a_card=deployed.a2a_card or {}, agent_name=self.agent_name, a2a_context_id=a2a_context_id, - artifact_id=self.resolve_artifact_id(context), + artifact_id=self._resolve_artifact_id(context), timeout_seconds=self.timeout_seconds, sandbox_type=deployed.sandbox_type, ) @@ -141,6 +141,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: "bundle_object_url": workspace.bundle_object_url, "source_agent_name": self.agent_name, "source_context_id": a2a_context_id, + "capture_step_id": self.id, }) # Per-MCP-service state dump. Only runs when env_id + diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index 22cb2644..6c160629 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +import re import pytest @@ -302,6 +303,7 @@ async def test_without_an_artifact_id_each_run_snapshots_into_its_own_artifact(c assert captured_into == ["inst-a__t-agent-snapshot", "inst-a__t-agent-snapshot", "inst-b__t-agent-snapshot"] assert first.metadata["agent_snapshots"][0]["id"] == "inst-a__t-agent-snapshot" + assert first.metadata["agent_snapshots"][0]["capture_step_id"] == "t-agent-snapshot" assert other.metadata["agent_snapshots"][0]["id"] == "inst-b__t-agent-snapshot" assert step.artifact_id is None @@ -316,7 +318,11 @@ async def test_an_explicit_artifact_id_is_used_for_every_run(captured_into): def test_a_local_instance_derives_a_local_artifact_id(): - assert _per_run_step().resolve_artifact_id(_run_ctx("@local/runs/r1")) == "@local/runs/r1__t-agent-snapshot" + assert _per_run_step()._resolve_artifact_id(_run_ctx("@local/runs/r1")) == "@local/runs/r1__t-agent-snapshot" + + +def test_without_an_instance_id_the_artifact_id_is_adhoc(): + assert re.fullmatch(r"adhoc-[0-9a-f]{12}__t-agent-snapshot", _per_run_step()._resolve_artifact_id(_run_ctx(None))) @pytest.mark.parametrize("artifact_id", [None, "shared-snapshot"]) From 2be4429fca088efe310010144f18d36b971bc42a Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 22 Oct 2026 18:02:41 -0700 Subject: [PATCH 3/5] feat(task-step): skip a per-run agent snapshot when the agent can't be snapshotted A per-run snapshot_agent_state (artifact_id omitted) is meant to be added to every run that prompts an agent, and not every agent serves the snapshot extension. Failing the step on each of their runs only adds a failed step to runs that have nothing to capture, so in that mode the step logs a warning, records no agent_snapshots entry and succeeds. A step with an explicit artifact_id still raises: a caller that names the artifact expects it written. A missing a2a_context_id still raises: the agent can be snapshotted, but the prompt it would be taken from did not record its conversation. FD-4006 Co-Authored-By: Claude Opus 5.5 --- .../task_steps/snapshot_agent_state.py | 11 +++++++- .../task_step/test_snapshot_agent_state.py | 27 +++++++++++++++++-- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/src/agent_env/task_step/task_steps/snapshot_agent_state.py b/src/agent_env/task_step/task_steps/snapshot_agent_state.py index eb4f91fa..8b7bf26c 100644 --- a/src/agent_env/task_step/task_steps/snapshot_agent_state.py +++ b/src/agent_env/task_step/task_steps/snapshot_agent_state.py @@ -6,6 +6,7 @@ import logging from typing import ClassVar, Optional +from agent_env.a2a_agent import A2AAgent from agent_env.store.ids import derive_id, is_local_id, validate_local_id from agent_env.task_step.context import TaskStepContext from agent_env.entity.refs import EntityRef, RefRole @@ -30,7 +31,9 @@ class SnapshotAgentStateTaskStep(TaskStep): Without an ``artifact_id`` the artifact is per run, derived from the instance id (``__``), so concurrent runs of one task never version each - other's snapshot and a retry of one instance reuses its id. + other's snapshot and a retry of one instance reuses its id. In that mode an agent + that does not advertise the snapshot extension is skipped, recording no entry, + rather than failing the step; with an explicit ``artifact_id`` it still raises. When ``env_id`` is set, the step ALSO reads each MCP service's current state from the deployed env (its ``data/get`` JSON, else its @@ -113,6 +116,12 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: a2a_url = deployed.a2a_url or deployed.api_url if not a2a_url: raise RuntimeError(f"Deployed agent '{self.agent_name}' has no a2a_url") + if not self.artifact_id and not A2AAgent.find_extension(deployed.a2a_card or {}, A2AAgent.EXT_SNAPSHOT): + logger.warning( + "snapshot_agent_state %s: agent '%s' does not advertise the snapshot extension; " + "skipping the per-run capture", self.id, self.agent_name, + ) + return context prompt = next((p for p in context.prompt_responses if p.prompt_id == self.prompt_id), None) if prompt is None: diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index 6c160629..ac898e09 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -4,12 +4,14 @@ from __future__ import annotations import json +import logging import re import pytest import agent_env.artifact as artifact_mod import agent_env.config as config_mod +from agent_env.a2a_agent import A2AAgent from agent_env.env import legacy_protocol from agent_env.env.env import DeployedGatewayEnv from agent_env.store import set_object_store @@ -272,9 +274,12 @@ def _per_run_step(artifact_id=None): ) -def _run_ctx(instance_id): +_SNAPSHOT_CARD = {"capabilities": {"extensions": [{"uri": A2AAgent.EXT_SNAPSHOT}]}} + + +def _run_ctx(instance_id, a2a_card=_SNAPSHOT_CARD): ctx = TaskStepContext(instance_id=instance_id) - ctx.deployed_agents = [DeployedAgent(agent_name="agent", api_url="https://agent")] + ctx.deployed_agents = [DeployedAgent(agent_name="agent", api_url="https://agent", a2a_card=a2a_card)] ctx.prompt_responses = [PromptResponse(prompt_id="main", response="ok", a2a_context_id="ctx-1")] return ctx @@ -317,6 +322,24 @@ async def test_an_explicit_artifact_id_is_used_for_every_run(captured_into): assert captured_into == ["shared-snapshot", "shared-snapshot"] +@pytest.mark.asyncio +async def test_a_per_run_capture_skips_an_agent_without_the_snapshot_extension(captured_into, caplog): + ctx = _run_ctx("inst-a", a2a_card={"capabilities": {"extensions": []}}) + with caplog.at_level(logging.WARNING, logger=mod.__name__): + assert await _per_run_step().execute(ctx) is ctx + + assert captured_into == [] + assert "agent_snapshots" not in ctx.metadata + assert "does not advertise the snapshot extension" in caplog.text + + +@pytest.mark.asyncio +async def test_an_explicit_artifact_id_still_fails_on_an_agent_without_the_snapshot_extension(): + ctx = _run_ctx("inst-a", a2a_card=None) + with pytest.raises(RuntimeError, match="does not advertise the snapshot extension"): + await _per_run_step("shared-snapshot").execute(ctx) + + def test_a_local_instance_derives_a_local_artifact_id(): assert _per_run_step()._resolve_artifact_id(_run_ctx("@local/runs/r1")) == "@local/runs/r1__t-agent-snapshot" From d49130586b6c2be47fa88108dbc9ff41b05d7934 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 22 Oct 2026 18:10:51 -0700 Subject: [PATCH 4/5] fix(task-step): skip a per-run agent snapshot only for a card without the extension An agent recorded without a card points at how the run context was built, not at an agent that can't be snapshotted, so it no longer skips: the step goes on to the capture and fails (non-fatally where configured) as before. FD-4006 Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/snapshot_agent_state.py | 11 ++++++++--- tst/unit/task_step/test_snapshot_agent_state.py | 6 ++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/agent_env/task_step/task_steps/snapshot_agent_state.py b/src/agent_env/task_step/task_steps/snapshot_agent_state.py index 8b7bf26c..8e726d71 100644 --- a/src/agent_env/task_step/task_steps/snapshot_agent_state.py +++ b/src/agent_env/task_step/task_steps/snapshot_agent_state.py @@ -32,8 +32,9 @@ class SnapshotAgentStateTaskStep(TaskStep): Without an ``artifact_id`` the artifact is per run, derived from the instance id (``__``), so concurrent runs of one task never version each other's snapshot and a retry of one instance reuses its id. In that mode an agent - that does not advertise the snapshot extension is skipped, recording no entry, - rather than failing the step; with an explicit ``artifact_id`` it still raises. + whose card does not advertise the snapshot extension is skipped, recording no entry, + rather than failing the step. With an explicit ``artifact_id``, or an agent recorded + without a card, it still raises. When ``env_id`` is set, the step ALSO reads each MCP service's current state from the deployed env (its ``data/get`` JSON, else its @@ -116,7 +117,11 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: a2a_url = deployed.a2a_url or deployed.api_url if not a2a_url: raise RuntimeError(f"Deployed agent '{self.agent_name}' has no a2a_url") - if not self.artifact_id and not A2AAgent.find_extension(deployed.a2a_card or {}, A2AAgent.EXT_SNAPSHOT): + if ( + not self.artifact_id + and deployed.a2a_card is not None + and not A2AAgent.find_extension(deployed.a2a_card, A2AAgent.EXT_SNAPSHOT) + ): logger.warning( "snapshot_agent_state %s: agent '%s' does not advertise the snapshot extension; " "skipping the per-run capture", self.id, self.agent_name, diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index ac898e09..9a36608a 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -333,6 +333,12 @@ async def test_a_per_run_capture_skips_an_agent_without_the_snapshot_extension(c assert "does not advertise the snapshot extension" in caplog.text +@pytest.mark.asyncio +async def test_a_per_run_capture_still_fails_on_an_agent_recorded_without_a_card(): + with pytest.raises(RuntimeError, match="does not advertise the snapshot extension"): + await _per_run_step().execute(_run_ctx("inst-a", a2a_card=None)) + + @pytest.mark.asyncio async def test_an_explicit_artifact_id_still_fails_on_an_agent_without_the_snapshot_extension(): ctx = _run_ctx("inst-a", a2a_card=None) From e030de2b41ebcf255f52023ec874a25e5f002a10 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Fri, 23 Oct 2026 09:30:51 -0700 Subject: [PATCH 5/5] fix(task-step): keep snapshot_agent_state's skip, id and CLI output to its own capture - The per-run skip no longer covers a step that also dumps per-service state (env_id and universe_artifact_id): the dump is written beside the workspace capture, so skipping it silently left the run without that state. Such a step still raises. The skip and the dump read one predicate, so they can't drift. - The derived artifact id is `__-agent-state`, so it can't share an id with the ids a prompt step's snapshot series derives from its own step id (`-workspace`, `snapshot-`). - `agent-env task run` shows a snapshot_agent_state step the row it wrote (matched on capture_step_id) and nothing when it wrote none, rather than whichever agent_snapshots row came last. - to_dict notes why it writes artifact_id even when None. FD-4006 Co-Authored-By: Claude Opus 5.5 --- src/agent_env/cli/task/run.py | 5 ++-- .../task_steps/snapshot_agent_state.py | 18 ++++++++---- tst/unit/cli/snapshot_step_output_test.py | 28 +++++++++++++++++++ .../task_step/test_snapshot_agent_state.py | 20 +++++++++---- 4 files changed, 59 insertions(+), 12 deletions(-) create mode 100644 tst/unit/cli/snapshot_step_output_test.py diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 3fceee14..dce593ff 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -172,8 +172,9 @@ def _step_output_lines(step, context) -> list[str]: return [_format_prompt_response(context.prompt_responses[-1])] if step.type == "snapshot_agent_state": snapshots = context.metadata.get("agent_snapshots") or [] - if snapshots: - return [_format_snapshot_capture(snapshots[-1])] + match = next((s for s in reversed(snapshots) if s.get("capture_step_id") == step.id), None) + if match: + return [_format_snapshot_capture(match)] if step.type == "modify_env_tool_access": changes = context.metadata.get("tool_access_changes") or [] match = next((c for c in reversed(changes) if c.get("step_id") == step.id), None) diff --git a/src/agent_env/task_step/task_steps/snapshot_agent_state.py b/src/agent_env/task_step/task_steps/snapshot_agent_state.py index 8e726d71..b491029d 100644 --- a/src/agent_env/task_step/task_steps/snapshot_agent_state.py +++ b/src/agent_env/task_step/task_steps/snapshot_agent_state.py @@ -30,11 +30,12 @@ class SnapshotAgentStateTaskStep(TaskStep): ``context.metadata['agent_snapshots']`` (a list of such entries). Without an ``artifact_id`` the artifact is per run, derived from the instance id - (``__``), so concurrent runs of one task never version each + (``__-agent-state``), so concurrent runs of one task never version each other's snapshot and a retry of one instance reuses its id. In that mode an agent whose card does not advertise the snapshot extension is skipped, recording no entry, - rather than failing the step. With an explicit ``artifact_id``, or an agent recorded - without a card, it still raises. + rather than failing the step. With an explicit ``artifact_id``, an agent recorded + without a card, or a per-service state dump configured (it is written beside the + workspace capture, so it cannot run without one), it still raises. When ``env_id`` is set, the step ALSO reads each MCP service's current state from the deployed env (its ``data/get`` JSON, else its @@ -82,6 +83,8 @@ def __init__( def to_dict(self) -> dict: base = super().to_dict() + # Written even when None: an older release then loads the step and fails only it at + # run time, where a missing key would make it refuse the whole task. base["artifact_id"] = self.artifact_id base["prompt_id"] = self.prompt_id base["agent_name"] = self.agent_name @@ -102,10 +105,14 @@ def from_dict(cls, data: dict) -> SnapshotAgentStateTaskStep: timeout_seconds=data.get("timeout_seconds", cls.DEFAULT_TIMEOUT_SECONDS), ) + @property + def _dumps_service_state(self) -> bool: + return bool(self.env_id and self.universe_artifact_id) + def _resolve_artifact_id(self, context: TaskStepContext) -> str: if self.artifact_id: return self.artifact_id - artifact_id = derive_id(rollout_base(self.id, context.instance_id), self.id) + artifact_id = derive_id(rollout_base(self.id, context.instance_id), f"{self.id}-agent-state") if is_local_id(artifact_id): validate_local_id(artifact_id) return artifact_id @@ -119,6 +126,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise RuntimeError(f"Deployed agent '{self.agent_name}' has no a2a_url") if ( not self.artifact_id + and not self._dumps_service_state and deployed.a2a_card is not None and not A2AAgent.find_extension(deployed.a2a_card, A2AAgent.EXT_SNAPSHOT) ): @@ -165,7 +173,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # JSON-stringified `{service: object_url}` (unsigned object-store references). # A downstream verifier re-signs each value and hands the map to its test # script, which accepts either a single URL or this map. - if self.env_id and self.universe_artifact_id: + if self._dumps_service_state: await self._capture_universe_state(context, workspace.capture_prefix) return context diff --git a/tst/unit/cli/snapshot_step_output_test.py b/tst/unit/cli/snapshot_step_output_test.py new file mode 100644 index 00000000..468829e7 --- /dev/null +++ b/tst/unit/cli/snapshot_step_output_test.py @@ -0,0 +1,28 @@ +"""`task run` shows a snapshot_agent_state step its own capture, never another writer's row.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from agent_env.cli.task.run import _step_output_lines +from agent_env.task_step.context import TaskStepContext + +_STEP = SimpleNamespace(type="snapshot_agent_state", id="capture") + + +def _context(*rows: dict) -> TaskStepContext: + return TaskStepContext(metadata={"agent_snapshots": list(rows)}) + + +def test_the_steps_own_capture_is_shown_over_a_later_row(): + lines = _step_output_lines(_STEP, _context( + {"id": "own", "version": 2, "capture_step_id": "capture"}, + {"id": "series-row", "version": 5, "capture_step_id": "ask"}, + )) + + assert len(lines) == 1 + assert "artifact_id: own" in lines[0] + + +def test_nothing_is_shown_without_a_row_from_the_step(): + assert _step_output_lines(_STEP, _context({"id": "series-row", "capture_step_id": "ask"}, {"id": "untagged"})) == [] diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index 9a36608a..8682c673 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -306,10 +306,10 @@ async def test_without_an_artifact_id_each_run_snapshots_into_its_own_artifact(c for ctx in (first, retry, other): await step.execute(ctx) - assert captured_into == ["inst-a__t-agent-snapshot", "inst-a__t-agent-snapshot", "inst-b__t-agent-snapshot"] - assert first.metadata["agent_snapshots"][0]["id"] == "inst-a__t-agent-snapshot" + assert captured_into == ["inst-a__t-agent-snapshot-agent-state", "inst-a__t-agent-snapshot-agent-state", "inst-b__t-agent-snapshot-agent-state"] + assert first.metadata["agent_snapshots"][0]["id"] == "inst-a__t-agent-snapshot-agent-state" assert first.metadata["agent_snapshots"][0]["capture_step_id"] == "t-agent-snapshot" - assert other.metadata["agent_snapshots"][0]["id"] == "inst-b__t-agent-snapshot" + assert other.metadata["agent_snapshots"][0]["id"] == "inst-b__t-agent-snapshot-agent-state" assert step.artifact_id is None @@ -333,6 +333,16 @@ async def test_a_per_run_capture_skips_an_agent_without_the_snapshot_extension(c assert "does not advertise the snapshot extension" in caplog.text +@pytest.mark.asyncio +async def test_a_per_run_capture_with_a_service_state_dump_still_fails_on_an_agent_without_the_extension(): + step = mod.SnapshotAgentStateTaskStep( + id="t-agent-snapshot", version=None, artifact_id=None, prompt_id="main", agent_name="agent", + env_id="env-1", universe_artifact_id="uni-1", + ) + with pytest.raises(RuntimeError, match="does not advertise the snapshot extension"): + await step.execute(_run_ctx("inst-a", a2a_card={"capabilities": {"extensions": []}})) + + @pytest.mark.asyncio async def test_a_per_run_capture_still_fails_on_an_agent_recorded_without_a_card(): with pytest.raises(RuntimeError, match="does not advertise the snapshot extension"): @@ -347,11 +357,11 @@ async def test_an_explicit_artifact_id_still_fails_on_an_agent_without_the_snaps def test_a_local_instance_derives_a_local_artifact_id(): - assert _per_run_step()._resolve_artifact_id(_run_ctx("@local/runs/r1")) == "@local/runs/r1__t-agent-snapshot" + assert _per_run_step()._resolve_artifact_id(_run_ctx("@local/runs/r1")) == "@local/runs/r1__t-agent-snapshot-agent-state" def test_without_an_instance_id_the_artifact_id_is_adhoc(): - assert re.fullmatch(r"adhoc-[0-9a-f]{12}__t-agent-snapshot", _per_run_step()._resolve_artifact_id(_run_ctx(None))) + assert re.fullmatch(r"adhoc-[0-9a-f]{12}__t-agent-snapshot-agent-state", _per_run_step()._resolve_artifact_id(_run_ctx(None))) @pytest.mark.parametrize("artifact_id", [None, "shared-snapshot"])