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/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..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 @@ -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 + (``__-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,7 +97,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 +105,18 @@ 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), 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: @@ -99,6 +124,17 @@ 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 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 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 08a46d70..8682c673 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/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