feat(task-step): derive snapshot_agent_state's artifact id per run when omitted

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)

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"))

src/agent_env/task_step/task_steps/snapshot_agent_state.py

@@ -6,10 +6,13 @@
 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
 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__)
 
@@ -23,9 +26,17 @@ 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
+    (``<instance_id>__<step id>-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``, 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
     ``/export-state``), uploads each as JSON alongside the workspace tarball, and
@@ -49,7 +60,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,
@@ -72,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
@@ -84,21 +97,44 @@ 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"),
             universe_artifact_id=data.get("universe_artifact_id"),
             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), f"{self.id}-agent-state")
+        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:
             raise RuntimeError(f"No deployed agent named '{self.agent_name}' found in context")
         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 self._dumps_service_state
+            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,
+            )
+            return context
 
         prompt = next((p for p in context.prompt_responses if p.prompt_id == self.prompt_id), None)
         if prompt is None:
@@ -115,7 +151,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,
         )
@@ -127,6 +163,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 +
@@ -136,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

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"})) == []

tst/unit/task_step/test_snapshot_agent_state.py

@@ -1,18 +1,23 @@
-"""``_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
+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
 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 +266,115 @@ 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",
+    )
+
+
+_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", a2a_card=a2a_card)]
+    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-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-agent-state"
+    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"]
+
+
+@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_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"):
+        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)
+    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-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-agent-state", _per_run_step()._resolve_artifact_id(_run_ctx(None)))
+
+
+@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