From 7f79286333f79a657fa356656777cbb8e7f06e3f Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 7 Oct 2026 14:39:31 -0700 Subject: [PATCH 01/12] feat(protocol): live trajectory log and cursor reads for trajectory/v1 Agents can now push trajectory events into a framework-owned, append-only log (`request.trajectory.append(event)`), and clients can read it while the task runs. `enable(TRAJECTORY_V1, live=True)` advertises two framework-owned cursor reads on the existing `get`: `{task_id, after, limit?}` and `{context_id, after, limit?}`, answered with `{context_id, task_id, state, format, events, next, has_more}`. - Additive: the `{task_id}`, `{task_id, objects}`, `{context_id}` and `{context_id, objects}` reads keep their request and response shapes. A `native_trajectory` is still the final record and is encoded as before, so the final object is unchanged for every existing agent. - When an agent appends and returns no `native_trajectory`, the log is the task's final trajectory. - States: pending, running, completed, failed, canceled. A terminal state is reported only once the log is sealed, and the log is sealed before the task's terminal A2A event. - Tasks are registered before the executor's first await, so an id returned by `message/send` is never 404. Context reads follow execution order. - With `live=True` the framework answers the `{context_id}` reads from its log unless the agent binds its own handlers, which still take precedence. - The completed-trajectory cache is replaced by a store bounded at 256 MiB and 1,024 contexts. Whole contexts are evicted, least recently used first, and never one with a running task. FD-3719 Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 84 ++- .../agentenv_protocol/a2a_agent/__init__.py | 10 + .../a2a_agent/_trajectory_log.py | 249 +++++++++ .../agentenv_protocol/a2a_agent/extensions.py | 79 ++- .../agentenv_protocol/a2a_agent/framework.py | 146 +++-- .../a2a_agent/tasks/__init__.py | 2 + .../agentenv_protocol/a2a_agent/tasks/v1.py | 63 ++- .../agentenv-protocol/tests/test_a2a_agent.py | 33 +- .../tests/test_a2a_object_transfers.py | 9 +- .../tests/test_a2a_trajectory_log.py | 502 ++++++++++++++++++ 10 files changed, 1100 insertions(+), 77 deletions(-) create mode 100644 packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py create mode 100644 packages/agentenv-protocol/tests/test_a2a_trajectory_log.py diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index cd84562d..37beb205 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -280,8 +280,9 @@ executor maps them to and from the wire-level `a2a.types.Task` lifecycle. `TaskRequest` is a frozen record whose nested JSON values are detached copies. Its `config`, `mcp_servers`, `skills`, `metadata`, and inbound `DataPart.data` retain their declared `dict`/`list` types, so normal Pydantic serialization, -copying, and `json.dumps(...)` work. Within `tasks.v1`, new request fields are -additive and have framework defaults. Agent code returns a `TaskResult` through +copying, and `json.dumps(...)` work. `TaskRequest.trajectory` is the exception: +it is the task's live trajectory log, described below. Within `tasks.v1`, new +request fields are additive and have framework defaults. Agent code returns a `TaskResult` through its factories or builder so additions to the result contract do not break existing handlers. @@ -304,7 +305,74 @@ The SDK does not retry tasks. A client that gives up on a task sends `tasks/cancel`. The framework marks the task canceled and cancels `run()`, which gets `asyncio.CancelledError` at its next `await`. A process `run()` started keeps running unless `run()` stops it, -so kill it before re-raising. +so kill it before re-raising. Trajectory events it appended, including any it +appends while cleaning up, are kept, and its log ends `canceled`. + +#### Live trajectories + +A runtime that sees its native events as they happen appends each complete +event to `request.trajectory` and enables `TRAJECTORY_V1` with `live=True`: + +```python +@a2a_agent(identity=..., extensions=(enable(TRAJECTORY_V1, live=True),)) +class MyAgent(AgentEnvAgent): + async def run(self, request: TaskRequest[MyAgentConfig]) -> TaskResult: + request.trajectory.set_format("my-cli-events/1") + execution = await run_my_cli(request, on_event=request.trajectory.append) + return TaskResult.text(execution.output) +``` + +`append(event)` takes a JSON object and serializes it at once, so an event +never changes after it is appended and a read never returns part of one. It does +no I/O; call it from one thread at a time. After the task ends the log refuses +further events. `set_format` names the events' versioned format, which readers +use to parse them. + +The framework keeps each task's log and serves it while the task runs. When the +result carries no `native_trajectory`, the log is the task's final trajectory: +the completed-task reads `{task_id}` and `{task_id, objects}` return it as a JSON +array. A `native_trajectory` remains the final record of a runtime that can only +produce one at the end, and takes precedence over appended events. + +`live=True` adds two cursor reads to `TRAJECTORY_V1.get`, both framework-owned: + +| Request | Reads | +| --- | --- | +| `{task_id, after, limit?}` | one task's events | +| `{context_id, after, limit?}` | the events of every task in a context, in execution order | + +Both answer `{context_id, task_id, state, format, events, next, has_more}` +(`TrajectoryEventsResponse`): + +- `events` are the events from position `after`, zero-based. `next` is + `after` plus the number returned, the position to read from next; `has_more` + says more events are already readable past this page. A page holds at most + 1,000 events and about 4 MiB, but always at least one event when one exists; + `limit` lowers the event count. `after` past the last event is HTTP 400. +- `state` is `pending` before the first event, `running` while the task runs, + then `completed`, `failed` or `canceled`. A terminal state is reported only + once the log can no longer change, so a reader that has reached `has_more: + false` in a terminal state has every event. The terminal state is recorded + before the task's terminal A2A status, so the completed-task reads succeed as + soon as a client sees the task end. +- A context read reports the context's current task, the running one or else the + latest, with its `format`; its `state` is `running` while any task runs. +- A task or context the agent has not seen, or no longer keeps, is HTTP 404. A + task is known from the moment `message/send` returns its id, so a 404 for an + id that was readable means it was evicted, and any cursor held for it, or for + its context, is no longer valid. + +`live=True` also enables the context reads `{context_id}` and +`{context_id, objects}`; the framework answers them with the context's events +unless the agent binds its own handlers, which take precedence. An agent that +overrides the whole `get` operation must accept the cursor requests too. +Clients detect live support from the card, with +`card_request_accepts(get_request, {"task_id", "after"})`. + +The logs are kept in memory, up to 256 MiB and 1,024 contexts. Past that, whole +contexts are dropped, least recently used first; a context with a task still +pending or running is never dropped. A running task's log is logged as a warning +once it passes 100 MiB. `enable(..., description="...")` is reserved for declarations carrying configuration or metadata. It preserves the agent-specific extension prose @@ -319,7 +387,8 @@ entry in `extensions=` is required. Configuration keywords are definition-owned rather than hardcoded in `enable()`. A configurable `ExtensionDefinition` supplies a named keyword-only `configuration_validator` returning `ExtensionConfiguration` with card -`wire_params`, internal `options`, and optional `features`. `enable()` only +`wire_params`, internal `options`, and optional `features` and request +`variants`. `enable()` only dispatches to that callable. Definitions without a validator reject configuration keywords, and third-party definitions use the same public API as the built-ins. @@ -416,9 +485,9 @@ async def get_live_trajectory(self, request: ContextObjectTrajectoryRequest): The last handler opts that agent into the optional live-context variant of `TRAJECTORY_V1.get` that uploads through an object grant; `TRAJECTORY_V1.get.context` (`ContextTrajectoryRequest`) is its inline -counterpart. Without them the generated card advertises only the -framework-owned completed-task variants: `{task_id}`, answered inline, and -`{task_id, objects}`, answered by upload. +counterpart. Without them, or `enable(TRAJECTORY_V1, live=True)`, the +generated card advertises only the framework-owned completed-task variants: +`{task_id}`, answered inline, and `{task_id, objects}`, answered by upload. Snapshot `save` and `load` are an atomic core contract, while its optional changelog handlers are enabled as an atomic feature group. Each operation has at most one response model. The framework validates a @@ -556,6 +625,7 @@ optional field, `|` an alternative request): | skill `add` | `{name, description, skill_md}` \| `{name, description, skill_bundle: {max_total_bytes, files: [{path, object: ReadObject}]}}` | `{name}` | | trajectory `get` | `{task_id}` \| `{context_id}` | `{trajectory}` | | | `{task_id \| context_id, objects: {trajectory: WriteObject}}` | `{objects: {trajectory: Uploaded}}` | +| | `{task_id \| context_id, after, limit?}` | `{context_id, task_id, state, format, events, next, has_more}` | | snapshot `save` | `{context_id, objects: {trajectory: WriteObject, workspace?: WriteObject}}` | `{context_id, objects: {trajectory: Uploaded, workspace?: Uploaded}}` | | snapshot `load` | `{objects: {trajectory: ReadObject, workspace?: ReadObject}, target_context_id?}` | `{context_id}` | | `enable-changelog` | `{write_namespace: WriteNamespaceGrant, roots?}` | `{roots}` | diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py index 1e978a33..93141797 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py @@ -16,6 +16,7 @@ AttributionProbeResponse, BundleSkillRequest, ChangelogIncrement, + ContextEventsTrajectoryRequest, ContextObjectTrajectoryRequest, ContextTrajectoryRequest, ExtensionActivation, @@ -48,9 +49,12 @@ SnapshotReadObjects, SnapshotUploadedObjects, SnapshotWriteObjects, + TaskEventsTrajectoryRequest, TaskObjectTrajectoryRequest, TaskTrajectoryRequest, + TrajectoryEventsResponse, TrajectoryObjectsResponse, + TrajectoryState, TrajectoryUploadedObjects, TrajectoryWriteObjects, TriggerDecideRequest, @@ -92,6 +96,7 @@ TaskResultBuilder, TaskStreamItem, TextPart, + TrajectoryLog, Usage, WriteOnly, ) @@ -131,6 +136,7 @@ "AttributionProbeResponse", "BundleSkillRequest", "ChangelogIncrement", + "ContextEventsTrajectoryRequest", "ContextObjectTrajectoryRequest", "ContextTrajectoryRequest", "DataPart", @@ -177,6 +183,7 @@ "SnapshotWriteObjects", "StagingStore", "TaskError", + "TaskEventsTrajectoryRequest", "TaskObjectTrajectoryRequest", "TaskOutcome", "TaskPart", @@ -187,7 +194,10 @@ "TaskStreamItem", "TaskTrajectoryRequest", "TextPart", + "TrajectoryEventsResponse", + "TrajectoryLog", "TrajectoryObjectsResponse", + "TrajectoryState", "TrajectoryUploadedObjects", "TrajectoryWriteObjects", "TransferError", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py new file mode 100644 index 00000000..6905eaa4 --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py @@ -0,0 +1,249 @@ +"""The framework's in-process store of task trajectories, served by the trajectory extension.""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from dataclasses import dataclass + +from .extensions import TrajectoryState +from .tasks.v1 import TrajectoryLog + +logger = logging.getLogger(__name__) + +MAX_CACHED_TRAJECTORY_BYTES = 256 * 1024 * 1024 +MAX_CACHED_CONTEXTS = 1_024 +MAX_TRAJECTORY_READ_EVENTS = 1_000 +MAX_TRAJECTORY_READ_BYTES = 4 * 1024 * 1024 + + +@dataclass(slots=True) +class _TaskRecord: + context_id: str + log: TrajectoryLog + started: bool = False + cancel_requested: bool = False + final_state: TrajectoryState | None = None + native: bytes | None = None + + @property + def state(self) -> TrajectoryState: + if self.final_state is not None: + return self.final_state + return TrajectoryState.RUNNING if len(self.log) else TrajectoryState.PENDING + + @property + def size_bytes(self) -> int: + return self.log._size_bytes + len(self.native or b"") + + +@dataclass(frozen=True, slots=True) +class TrajectoryPage: + context_id: str + task_id: str + state: TrajectoryState + format: str | None + events: list[bytes] + total: int + + +class TaskTrajectories: + """Each task's log and final record, grouped by context in execution order. + + A whole context is the unit of eviction, least recently used first, so positions in a + context read never shift; a context with a task pending or running is never evicted.""" + + def __init__( + self, + max_bytes: int = MAX_CACHED_TRAJECTORY_BYTES, + max_contexts: int = MAX_CACHED_CONTEXTS, + ) -> None: + if max_bytes < 1 or max_contexts < 1: + raise ValueError("trajectory cache bounds must be positive") + self._max_bytes = max_bytes + self._max_contexts = max_contexts + self._tasks: dict[str, _TaskRecord] = {} + self._contexts: OrderedDict[str, list[str]] = OrderedDict() + + def register(self, task_id: str, context_id: str) -> TrajectoryLog: + """The task's log, created ``pending`` the first time the task is seen.""" + record = self._tasks.get(task_id) + if record is None: + log = TrajectoryLog() + log._task_id = task_id + record = self._tasks[task_id] = _TaskRecord(context_id=context_id, log=log) + self._contexts.setdefault(context_id, []).append(task_id) + self._touch(record.context_id) + return record.log + + def start(self, task_id: str) -> None: + """Place the task after the context's started tasks, as it begins to run. + + Only tasks that have not started move; they have no events, so no position shifts.""" + record = self._tasks[task_id] + if record.started: + return + order = self._contexts[record.context_id] + order.remove(task_id) + started = sum(1 for other in order if self._tasks[other].started) + order.insert(started, task_id) + record.started = True + self._touch(record.context_id) + + def request_cancel(self, task_id: str) -> None: + """Make ``canceled`` the outcome the task is sealed with, whatever ``run`` returns.""" + record = self._tasks.get(task_id) + if record is not None: + record.cancel_requested = True + + def seal( + self, task_id: str, state: TrajectoryState, *, native: bytes | None = None + ) -> None: + """End the task's trajectory; the first seal wins.""" + record = self._tasks[task_id] + if record.final_state is not None: + return + if native is not None: + if len(record.log): + logger.warning( + "task %s appended trajectory events and returned a native trajectory; " + "the native trajectory is its final record", + task_id, + ) + record.native = native + record.log._sealed = True + record.final_state = TrajectoryState.CANCELED if record.cancel_requested else state + self._evict() + + def final(self, task_id: str) -> bytes | None: + """The ended task's final trajectory as a JSON array, or None if there is none yet.""" + record = self._tasks.get(task_id) + if record is None or record.final_state is None: + return None + self._touch(record.context_id) + if record.native is not None: + return record.native + if not len(record.log): + return None + return _json_array(record.log._events) + + def context_events(self, context_id: str) -> bytes | None: + """Every event of the context's tasks in order, as a JSON array.""" + records = self._context_records(context_id) + if records is None: + return None + return _json_array([event for record in records for event in record.log._events]) + + def task_page( + self, task_id: str, after: int, max_events: int, max_bytes: int + ) -> TrajectoryPage | None: + record = self._tasks.get(task_id) + if record is None: + return None + self._touch(record.context_id) + # State before length: a terminal state means the length can no longer grow. + state = record.state + events = record.log._events + total = len(events) + return TrajectoryPage( + context_id=record.context_id, + task_id=task_id, + state=state, + format=record.log.format, + events=_page(events, after, total, max_events, max_bytes), + total=total, + ) + + def context_page( + self, context_id: str, after: int, max_events: int, max_bytes: int + ) -> TrajectoryPage | None: + records = self._context_records(context_id) + if records is None: + return None + task_ids = self._contexts[context_id] + states = [record.state for record in records] + lengths = [len(record.log) for record in records] + current = next( + (index for index, record in enumerate(records) if record.final_state is None), + len(records) - 1, + ) + if TrajectoryState.RUNNING in states: + state = TrajectoryState.RUNNING + elif TrajectoryState.PENDING in states: + state = TrajectoryState.PENDING + else: + state = states[-1] + page: list[bytes] = [] + offset = 0 + for record, length in zip(records, lengths): + if len(page) >= max_events: + break + if after < offset + length: + start = max(after - offset, 0) + page.extend( + _page( + record.log._events, + start, + length, + max_events - len(page), + max_bytes - sum(map(len, page)), + first=not page, + ) + ) + offset += length + return TrajectoryPage( + context_id=context_id, + task_id=task_ids[current], + state=state, + format=records[current].log.format, + events=page, + total=sum(lengths), + ) + + def _context_records(self, context_id: str) -> list[_TaskRecord] | None: + task_ids = self._contexts.get(context_id) + if not task_ids: + return None + self._touch(context_id) + return [self._tasks[task_id] for task_id in task_ids] + + def _touch(self, context_id: str) -> None: + if context_id in self._contexts: + self._contexts.move_to_end(context_id) + + def _evict(self) -> None: + total = sum(record.size_bytes for record in self._tasks.values()) + for context_id in tuple(self._contexts): + if total <= self._max_bytes and len(self._contexts) <= self._max_contexts: + return + task_ids = self._contexts[context_id] + if any(self._tasks[task_id].final_state is None for task_id in task_ids): + continue + del self._contexts[context_id] + for task_id in task_ids: + total -= self._tasks.pop(task_id).size_bytes + + +def _page( + events: list[bytes], + start: int, + stop: int, + max_events: int, + max_bytes: int, + *, + first: bool = True, +) -> list[bytes]: + """Events from ``start``, within both budgets; the first event of a read is always + returned, so one larger than ``max_bytes`` cannot stall a reader.""" + page: list[bytes] = [] + used = 0 + for encoded in events[start:min(stop, start + max_events)]: + if used + len(encoded) > max_bytes and (page or not first): + break + page.append(encoded) + used += len(encoded) + return page + + +def _json_array(events: list[bytes]) -> bytes: + return b"[" + b",".join(events) + b"]" diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py index bbe10112..0a049002 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py @@ -17,6 +17,7 @@ ConfigDict, Field, PositiveInt, + StrictInt, ValidationError, model_validator, ) @@ -132,6 +133,42 @@ class TrajectoryObjectsResponse(ExtensionResponse): objects: TrajectoryUploadedObjects +class TrajectoryState(str, Enum): + """Where a trajectory stands: ``pending`` has no events yet, ``running`` is still + gaining them, and the terminal states gain no more.""" + + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELED = "canceled" + + +class TaskEventsTrajectoryRequest(ExtensionRequest): + task_id: str + after: StrictInt = Field(ge=0) + limit: StrictInt | None = Field(default=None, ge=1) + + +class ContextEventsTrajectoryRequest(ExtensionRequest): + context_id: str + after: StrictInt = Field(ge=0) + limit: StrictInt | None = Field(default=None, ge=1) + + +class TrajectoryEventsResponse(ExtensionResponse): + """Events from position ``after`` on; ``next`` is the position to read from next, and + ``has_more`` says more events are already readable past this page.""" + + context_id: str + task_id: str + state: TrajectoryState + format: str | None = None + events: list[Any] + next: int + has_more: bool + + def _opaque_snapshot_objects( objects: SnapshotWriteObjects | SnapshotReadObjects, ) -> SnapshotWriteObjects | SnapshotReadObjects: @@ -573,6 +610,7 @@ class ExtensionConfiguration: wire_params: Mapping[str, Any] = field(default_factory=dict) options: Mapping[str, Any] = field(default_factory=dict) features: frozenset[str] = frozenset() + variants: Mapping[str, frozenset[str]] = field(default_factory=dict) def __post_init__(self) -> None: object.__setattr__( @@ -580,6 +618,13 @@ def __post_init__(self) -> None: ) object.__setattr__(self, "options", MappingProxyType(dict(self.options))) object.__setattr__(self, "features", frozenset(self.features)) + object.__setattr__( + self, + "variants", + MappingProxyType( + {operation: frozenset(names) for operation, names in self.variants.items()} + ), + ) @dataclass(frozen=True, slots=True) @@ -872,6 +917,18 @@ def _configure_install(**declaration: Any) -> ExtensionConfiguration: return ExtensionConfiguration(wire_params=declaration) +def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: + """``live=True`` advertises the cursor reads, which serve the framework's log while a + task runs, and the context reads the framework answers when the agent binds none.""" + if not isinstance(live, bool): + raise TypeError("live must be a boolean") + if not live: + return ExtensionConfiguration() + return ExtensionConfiguration( + variants={"get": {"task_events", "context_events", "context", "context_objects"}} + ) + + SDK = ImplementationOwner.SDK RUNTIME = ImplementationOwner.RUNTIME @@ -972,10 +1029,23 @@ def _configure_install(**declaration: Any) -> ExtensionConfiguration: support_required=False, implementation=RUNTIME, ), + RequestVariant( + "task_events", + TaskEventsTrajectoryRequest, + support_required=False, + implementation=SDK, + ), + RequestVariant( + "context_events", + ContextEventsTrajectoryRequest, + support_required=False, + implementation=SDK, + ), ), ), ) }, + configuration_validator=_configure_trajectory, ) SNAPSHOT_V1 = ExtensionDefinition( @@ -1158,14 +1228,17 @@ def enable( selected_features = set(features) selected_features.update(validated.features) + selected_variants = dict(validated.variants) + for operation, names in (variants or {}).items(): + selected_variants[operation] = selected_variants.get( + operation, frozenset() + ) | frozenset(names) return ExtensionActivation( definition=definition, description=description, features=frozenset(selected_features), - variants={ - operation: frozenset(names) for operation, names in (variants or {}).items() - }, + variants=selected_variants, wire_params=validated.wire_params, options=validated.options, required=required, diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py index 5c3c3436..84fd8a07 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py @@ -28,6 +28,11 @@ from starlette.routing import Route from starlette.types import ASGIApp, Receive, Scope, Send +from ._trajectory_log import ( + MAX_TRAJECTORY_READ_BYTES, + MAX_TRAJECTORY_READ_EVENTS, + TaskTrajectories, +) from ._triggers import TriggerEngine, TriggerError from .extensions import ( AGENT_CONFIG_V1, @@ -35,13 +40,19 @@ SKILL_CONFIG_V1, TRAJECTORY_V1, TRIGGERS_V1, + ContextEventsTrajectoryRequest, + ContextObjectTrajectoryRequest, + ContextTrajectoryRequest, ExtensionActivation, ExtensionDefinition, ImplementationOwner, McpAddRequest, OperationReference, + TaskEventsTrajectoryRequest, TaskObjectTrajectoryRequest, TaskTrajectoryRequest, + TrajectoryEventsResponse, + TrajectoryState, TriggerDecideRequest, TriggerRegisterRequest, enable, @@ -89,7 +100,6 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: _AGENT_DEFINITION = "_agentenv_a2a_definition" _AgentT = TypeVar("_AgentT", bound="AgentEnvAgent") _MAX_CACHED_CONTEXTS = 1_024 -_MAX_CACHED_TRAJECTORIES = 1_024 def _config_wire_names(config: type[AgentConfig]) -> dict[str, str]: @@ -155,36 +165,6 @@ def set(self, context_id: str, session_ref: str) -> None: self._entries.popitem(last=False) -class _BoundedTaskTrajectories: - """Least-recently-used cache of completed native task trajectories.""" - - def __init__(self, max_entries: int = _MAX_CACHED_TRAJECTORIES) -> None: - if max_entries < 1: - raise ValueError("max_entries must be positive") - self._max_entries = max_entries - self._entries: OrderedDict[str, Any] = OrderedDict() - - def __setitem__(self, task_id: str, trajectory: Any) -> None: - self._entries.pop(task_id, None) - self._entries[task_id] = trajectory - while len(self._entries) > self._max_entries: - self._entries.popitem(last=False) - - def get(self, task_id: str) -> Any | None: - try: - trajectory = self._entries.pop(task_id) - except KeyError: - return None - self._entries[task_id] = trajectory - return trajectory - - def pop(self, task_id: str, default: Any = None) -> Any: - return self._entries.pop(task_id, default) - - def __len__(self) -> int: - return len(self._entries) - - @dataclass(slots=True) class _ContextLockState: lock: asyncio.Lock = field(default_factory=asyncio.Lock) @@ -456,6 +436,10 @@ def _thaw(value: Any) -> Any: return value +def _encode_trajectory(payload: Any) -> bytes: + return json.dumps(_thaw(payload), separators=(",", ":"), default=str).encode() + + def _validation_error_detail(exc: ValidationError) -> str: """Return client-safe Pydantic errors without rejected values or docs URLs.""" return exc.json(include_input=False, include_url=False) @@ -557,7 +541,7 @@ def __init__( self.config_model ) self.mcp_servers: dict[str, dict[str, Any]] = {} - self.task_trajectories = _BoundedTaskTrajectories() + self.task_trajectories = TaskTrajectories() self._context_sessions = _BoundedContextSessions() self.skills: list[Mapping[str, Any]] = [] self.identity_skills: dict[str, str] = {} @@ -747,20 +731,35 @@ async def skill_list(self) -> dict[str, Any]: return {"skills": skills} async def trajectory_get( - self, request: TaskTrajectoryRequest | TaskObjectTrajectoryRequest - ) -> dict[str, Any]: - task_id = request.task_id - trajectory = self.task_trajectories.get(task_id) - if trajectory is None: - raise HTTPException( - status_code=404, - detail=f"No completed trajectory available for task '{task_id}'", - ) - if isinstance(request, TaskTrajectoryRequest): - return {"trajectory": _thaw(trajectory.payload)} - body = json.dumps( - _thaw(trajectory.payload), separators=(",", ":"), default=str - ).encode() + self, + request: TaskTrajectoryRequest + | TaskObjectTrajectoryRequest + | TaskEventsTrajectoryRequest + | ContextTrajectoryRequest + | ContextObjectTrajectoryRequest + | ContextEventsTrajectoryRequest, + ) -> Any: + if isinstance( + request, (TaskEventsTrajectoryRequest, ContextEventsTrajectoryRequest) + ): + return self._trajectory_events(request) + if isinstance(request, (TaskTrajectoryRequest, TaskObjectTrajectoryRequest)): + body = self.task_trajectories.final(request.task_id) + if body is None: + raise HTTPException( + status_code=404, + detail=( + f"No completed trajectory available for task '{request.task_id}'" + ), + ) + else: + body = self.task_trajectories.context_events(request.context_id) + if body is None: + raise HTTPException( + status_code=404, detail=f"Unknown context '{request.context_id}'" + ) + if isinstance(request, (TaskTrajectoryRequest, ContextTrajectoryRequest)): + return {"trajectory": json.loads(body)} uploaded = await upload(request.objects.trajectory, body) return { "objects": { @@ -768,6 +767,40 @@ async def trajectory_get( } } + def _trajectory_events( + self, request: TaskEventsTrajectoryRequest | ContextEventsTrajectoryRequest + ) -> TrajectoryEventsResponse: + max_events = min( + request.limit or MAX_TRAJECTORY_READ_EVENTS, MAX_TRAJECTORY_READ_EVENTS + ) + if isinstance(request, TaskEventsTrajectoryRequest): + page = self.task_trajectories.task_page( + request.task_id, request.after, max_events, MAX_TRAJECTORY_READ_BYTES + ) + unknown = f"Unknown task '{request.task_id}'" + else: + page = self.task_trajectories.context_page( + request.context_id, request.after, max_events, MAX_TRAJECTORY_READ_BYTES + ) + unknown = f"Unknown context '{request.context_id}'" + if page is None: + raise HTTPException(status_code=404, detail=unknown) + if request.after > page.total: + raise HTTPException( + status_code=400, + detail=f"after={request.after} is past the trajectory's {page.total} events", + ) + next_position = request.after + len(page.events) + return TrajectoryEventsResponse( + context_id=page.context_id, + task_id=page.task_id, + state=page.state, + format=page.format, + events=[json.loads(event) for event in page.events], + next=next_position, + has_more=next_position < page.total, + ) + async def triggers_register( self, request: TriggerRegisterRequest ) -> dict[str, Any]: @@ -1121,6 +1154,9 @@ async def execute(self, context: Any, event_queue: Any) -> None: raise ServerError(error=InvalidParamsError(message=str(exc))) from exc updater = TaskUpdater(event_queue, task.id, task.context_id) + trajectories = self._services.task_trajectories + # Before any await: once the client holds this task id, the log knows it. + trajectory = trajectories.register(task.id, task.context_id) try: if context.current_task is None: await event_queue.enqueue_event(task) @@ -1132,6 +1168,7 @@ async def execute(self, context: Any, event_queue: Any) -> None: **({"role": task_config.role} if task_config.role is not None else {}), } async with self._context_locks.acquire(task.context_id): + trajectories.start(task.id) request = TaskRequest( task_id=task.id, context_id=task.context_id, @@ -1144,6 +1181,7 @@ async def execute(self, context: Any, event_queue: Any) -> None: metadata=metadata, session_ref=self._services.session_ref_for_context(task.context_id), workspace=self._workspace, + trajectory=trajectory, ) if self._streaming: result = await self._run_streaming(request, updater, event_queue) @@ -1155,16 +1193,27 @@ async def execute(self, context: Any, event_queue: Any) -> None: self._services.set_session_ref_for_context( task.context_id, result.session_ref ) - if result.native_trajectory is not None: - self._services.task_trajectories[task.id] = result.native_trajectory parts = _to_a2a_parts(result) + # Sealed before the terminal event, so a client that sees the task end can read it. + native = result.native_trajectory + trajectories.seal( + task.id, + TrajectoryState.FAILED + if result.outcome is TaskOutcome.FAILED + else TrajectoryState.COMPLETED, + native=None if native is None else _encode_trajectory(native.payload), + ) message = updater.new_agent_message(parts=parts) if result.outcome is TaskOutcome.FAILED: await updater.failed(message=message) else: await updater.complete(message=message) + except asyncio.CancelledError: + trajectories.seal(task.id, TrajectoryState.CANCELED) + raise except Exception: + trajectories.seal(task.id, TrajectoryState.FAILED) correlation_id = uuid.uuid4().hex logger.exception( "unhandled task execution error (correlation_id=%s)", @@ -1250,6 +1299,7 @@ async def cancel(self, context: Any, event_queue: Any) -> None: ``CancelledError`` at its next ``await`` and should stop whatever it started, such as a CLI process.""" from a2a.server.tasks import TaskUpdater + self._services.task_trajectories.request_cancel(context.task_id) await TaskUpdater(event_queue, context.task_id, context.context_id).cancel() diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/__init__.py index e9700b19..2eae8ae6 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/__init__.py @@ -15,6 +15,7 @@ TaskResultBuilder, TaskStreamItem, TextPart, + TrajectoryLog, Usage, WriteOnly, ) @@ -34,6 +35,7 @@ "TaskResultBuilder", "TaskStreamItem", "TextPart", + "TrajectoryLog", "Usage", "WriteOnly", ] diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py index 139c6dd2..657121da 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py @@ -8,6 +8,8 @@ from __future__ import annotations +import json +import logging from collections.abc import AsyncIterator, Awaitable, Iterable, Mapping from dataclasses import dataclass, field from enum import Enum @@ -17,6 +19,9 @@ from a2a.types import TaskStatusUpdateEvent from pydantic import BaseModel, ConfigDict, Field +logger = logging.getLogger(__name__) + +_LARGE_LIVE_TRAJECTORY_BYTES = 100 * 1024 * 1024 _WriteOnlyT = TypeVar("_WriteOnlyT") WriteOnly: TypeAlias = Annotated[ @@ -125,9 +130,62 @@ class AgentConfig(BaseModel): ConfigT = TypeVar("ConfigT", bound=AgentConfig) +class TrajectoryLog: + """Append-only trajectory of one task, readable while the task runs. + + Each event is serialized as it is appended, so an event once read never changes; append + from one thread at a time. Unless the result carries a ``native_trajectory``, the log is + the task's final trajectory.""" + + __slots__ = ("_events", "_size_bytes", "_format", "_sealed", "_task_id", "_warned") + + def __init__(self) -> None: + self._events: list[bytes] = [] + self._size_bytes = 0 + self._format: str | None = None + self._sealed = False + self._task_id: str | None = None + self._warned = False + + @property + def format(self) -> str | None: + return self._format + + def set_format(self, format: str) -> None: + """Name the events' versioned format, such as ``my-cli-events/1``.""" + if not isinstance(format, str) or not format.strip(): + raise ValueError("trajectory format must be a non-empty string") + self._ensure_open() + self._format = format + + def append(self, event: Mapping[str, Any]) -> None: + """Record one complete event, a JSON object.""" + if not isinstance(event, Mapping): + raise TypeError("a trajectory event must be a JSON object") + self._ensure_open() + encoded = json.dumps(dict(event), separators=(",", ":"), allow_nan=False).encode() + self._events.append(encoded) + self._size_bytes += len(encoded) + if not self._warned and self._size_bytes > _LARGE_LIVE_TRAJECTORY_BYTES: + self._warned = True + logger.warning( + "trajectory of running task %s exceeds %d bytes", + self._task_id, + _LARGE_LIVE_TRAJECTORY_BYTES, + ) + + def __len__(self) -> int: + return len(self._events) + + def _ensure_open(self) -> None: + if self._sealed: + raise RuntimeError("the task has ended; its trajectory can no longer change") + + @dataclass(frozen=True, slots=True, kw_only=True) class TaskRequest(Generic[ConfigT]): - """Frozen request record with detached, JSON-native nested values.""" + """Frozen request record with detached, JSON-native nested values, plus the + task's live trajectory log.""" task_id: str context_id: str @@ -138,6 +196,9 @@ class TaskRequest(Generic[ConfigT]): metadata: Mapping[str, Any] = field(default_factory=dict) session_ref: str | None = None workspace: Path | None = None + trajectory: TrajectoryLog = field( + default_factory=TrajectoryLog, repr=False, compare=False + ) def __post_init__(self) -> None: object.__setattr__(self, "parts", tuple(self.parts)) diff --git a/packages/agentenv-protocol/tests/test_a2a_agent.py b/packages/agentenv-protocol/tests/test_a2a_agent.py index 51e6c594..92c7c4b6 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent.py @@ -72,6 +72,7 @@ TaskResult, TaskTrajectoryRequest, TextPart, + TrajectoryState, TriggerDecideRequest, TriggerRegisterRequest, Usage, @@ -90,10 +91,10 @@ TriggerError, TriggerTimeoutError, ) +from agentenv_protocol.a2a_agent._trajectory_log import TaskTrajectories from agentenv_protocol.a2a_agent.framework import ( _BoundedContextLocks, _BoundedContextSessions, - _BoundedTaskTrajectories, _SdkServices, _StandardExecutor, ) @@ -241,18 +242,21 @@ def test_context_session_cache_evicts_least_recently_used_context() -> None: assert sessions.get("third") == "session-3" -def test_trajectory_cache_evicts_least_recently_used_task() -> None: - trajectories = _BoundedTaskTrajectories(max_entries=2) - trajectories["first"] = {"event": 1} - trajectories["second"] = {"event": 2} - assert trajectories.get("first") == {"event": 1} +def test_trajectory_cache_evicts_least_recently_used_context() -> None: + trajectories = TaskTrajectories(max_contexts=2) + for number in (1, 2): + trajectories.register(f"task-{number}", f"context-{number}") + trajectories.seal( + f"task-{number}", TrajectoryState.COMPLETED, native=f"[{number}]".encode() + ) + assert trajectories.final("task-1") == b"[1]" - trajectories["third"] = {"event": 3} + trajectories.register("task-3", "context-3") + trajectories.seal("task-3", TrajectoryState.COMPLETED, native=b"[3]") - assert trajectories.get("first") == {"event": 1} - assert trajectories.get("second") is None - assert trajectories.get("third") == {"event": 3} - assert len(trajectories) == 2 + assert trajectories.final("task-1") == b"[1]" + assert trajectories.final("task-2") is None + assert trajectories.final("task-3") == b"[3]" @pytest.mark.asyncio @@ -2044,15 +2048,16 @@ async def upload_by_context(self, request: ContextObjectTrajectoryRequest): ], } - with pytest.raises(TypeError, match="unsupported configuration"): + with pytest.raises(TypeError, match="unexpected keyword argument 'source'"): enable(TRAJECTORY_V1, source="task_result") @pytest.mark.asyncio async def test_default_handler_delegation_accepts_sdk_request_variant() -> None: services = _SdkServices([enable(TRAJECTORY_V1)], None) - services.task_trajectories["task-1"] = NativeTrajectory( - format="events/v1", payload=[{"type": "result"}] + services.task_trajectories.register("task-1", "context-1") + services.task_trajectories.seal( + "task-1", TrajectoryState.COMPLETED, native=b'[{"type":"result"}]' ) result = await DefaultExtensionHandlers(services).call( diff --git a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py index d2bcdb99..20b53c4f 100644 --- a/packages/agentenv-protocol/tests/test_a2a_object_transfers.py +++ b/packages/agentenv-protocol/tests/test_a2a_object_transfers.py @@ -19,7 +19,6 @@ AgentEnvAgent, AgentIdentity, NamespaceChangelogEnableRequest, - NativeTrajectory, ObjectChangelogApplyRequest, ObjectSnapshotLoadRequest, ObjectSnapshotSaveRequest, @@ -32,6 +31,7 @@ TaskRequest, TaskResult, TaskTrajectoryRequest, + TrajectoryState, TrajectoryWriteObjects, a2a_agent, card_request_accepts, @@ -549,8 +549,9 @@ async def test_sdk_trajectory_handler_uploads_portable_object( monkeypatch: pytest.MonkeyPatch, ) -> None: services = _SdkServices((enable(TRAJECTORY_V1),), None) - services.task_trajectories["task-1"] = NativeTrajectory( - format="events/v1", payload=[{"type": "result"}] + services.task_trajectories.register("task-1", "context-1") + services.task_trajectories.seal( + "task-1", TrajectoryState.COMPLETED, native=b'[{"type":"result"}]' ) def handle(request: httpx.Request) -> httpx.Response: @@ -573,7 +574,7 @@ def handle(request: httpx.Request) -> httpx.Response: } } } - assert services.task_trajectories.get("task-1") is not None + assert services.task_trajectories.final("task-1") is not None @pytest.mark.asyncio diff --git a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py new file mode 100644 index 00000000..a948e69f --- /dev/null +++ b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py @@ -0,0 +1,502 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import threading +import time +from typing import Any, AsyncIterator + +import pytest +from starlette.testclient import TestClient + +from agentenv_protocol.a2a_agent import ( + TRAJECTORY_V1, + AgentEnvAgent, + AgentIdentity, + ContextTrajectoryRequest, + DefaultExtensionHandlers, + TaskEventsTrajectoryRequest, + TaskObjectTrajectoryRequest, + TaskProgress, + TaskRequest, + TaskResult, + TaskTrajectoryRequest, + TrajectoryLog, + TrajectoryState, + a2a_agent, + build_registry, + card_request_accepts, + enable, + extension, +) +from agentenv_protocol.a2a_agent._trajectory_log import TaskTrajectories +from agentenv_protocol.a2a_agent.framework import _SdkServices +from agentenv_protocol.a2a_agent.tasks import v1 as tasks_v1 + +_IDENTITY = AgentIdentity(name="live", description="test", version="1") + + +class _Gates: + """Lets a test hold the agent before it appends (``go``) and before it returns + (``release``); both are set by default.""" + + def __init__(self, *, go: bool = True, release: bool = True) -> None: + self.go = threading.Event() + self.release = threading.Event() + if go: + self.go.set() + if release: + self.release.set() + + +async def _wait(event: threading.Event) -> None: + while not event.is_set(): + await asyncio.sleep(0.01) + + +def _live_agent(gates: _Gates, *, events: int = 3, result: str = "done") -> AgentEnvAgent: + @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) + class Agent(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.set_format("test-events/1") + await _wait(gates.go) + for index in range(events): + request.trajectory.append({"text": request.parts[0].text, "index": index}) + await _wait(gates.release) + if result == "failed": + return TaskResult.failure("agent.failed", "it failed") + if result == "raise": + raise RuntimeError("boom") + return TaskResult.text(result) + + return Agent() + + +def _send(client: TestClient, text: str, message_id: str, context_id: str = "context-1") -> str: + response = client.post( + "/a2a", + json={ + "jsonrpc": "2.0", + "id": message_id, + "method": "message/send", + "params": { + "message": { + "kind": "message", + "messageId": message_id, + "role": "user", + "contextId": context_id, + "parts": [{"kind": "text", "text": text}], + }, + "configuration": {"blocking": False}, + }, + }, + ) + return response.json()["result"]["id"] + + +def _task_state(client: TestClient, task_id: str) -> str: + response = client.post( + "/a2a", + json={"jsonrpc": "2.0", "id": "get", "method": "tasks/get", "params": {"id": task_id}}, + ) + return response.json()["result"]["status"]["state"] + + +def _cancel(client: TestClient, task_id: str) -> dict[str, Any]: + return client.post( + "/a2a", + json={ + "jsonrpc": "2.0", + "id": "cancel", + "method": "tasks/cancel", + "params": {"id": task_id}, + }, + ).json() + + +def _trajectory(client: TestClient, payload: dict[str, Any]) -> Any: + return client.post("/ext/trajectory", json=payload) + + +def _within(seconds: float, condition) -> bool: + deadline = time.monotonic() + seconds + while not condition(): + if time.monotonic() > deadline: + return False + time.sleep(0.01) + return True + + +def _events(text: str, count: int) -> list[dict[str, Any]]: + return [{"text": text, "index": index} for index in range(count)] + + +def test_live_card_advertises_the_cursor_reads() -> None: + app = _live_agent(_Gates()).create_app() + card = app.state.agentenv_a2a.card.model_dump(mode="json", exclude_none=True) + (extension_card,) = [ + item for item in card["capabilities"]["extensions"] if item["uri"] == TRAJECTORY_V1.uri + ] + request = extension_card["params"]["methods"]["get"]["request"] + + assert request == { + "oneOf": [ + {"required": ["task_id"]}, + {"required": ["task_id", "objects"]}, + {"required": ["context_id"]}, + {"required": ["context_id", "objects"]}, + {"required": ["task_id", "after"], "optional": ["limit"]}, + {"required": ["context_id", "after"], "optional": ["limit"]}, + ] + } + assert card_request_accepts(request, {"task_id"}) + assert card_request_accepts(request, {"task_id", "after"}) + assert card_request_accepts(request, {"task_id", "after", "limit"}) + assert card_request_accepts(request, {"context_id", "after"}) + with pytest.raises(TypeError, match="live must be a boolean"): + enable(TRAJECTORY_V1, live="yes") + + +def test_a_full_get_override_must_take_the_cursor_reads_when_live() -> None: + class Agent: + @extension(TRAJECTORY_V1.get) + async def get(self, request: TaskTrajectoryRequest | TaskObjectTrajectoryRequest): + return {"trajectory": []} + + build_registry(Agent(), [enable(TRAJECTORY_V1)]) + with pytest.raises(ValueError, match="must annotate"): + build_registry(Agent(), [enable(TRAJECTORY_V1, live=True)]) + + +def test_task_read_follows_a_running_task() -> None: + gates = _Gates(go=False, release=False) + with TestClient(_live_agent(gates).create_app()) as client: + task_id = _send(client, "hello", "message-1") + pending = _trajectory(client, {"task_id": task_id, "after": 0}) + assert pending.status_code == 200 + assert {key: pending.json()[key] for key in ("task_id", "state", "events", "next")} == { + "task_id": task_id, + "state": "pending", + "events": [], + "next": 0, + } + + gates.go.set() + assert _within( + 5, + lambda: _trajectory(client, {"task_id": task_id, "after": 0}).json()["next"] == 3, + ) + first = _trajectory(client, {"task_id": task_id, "after": 0, "limit": 2}).json() + assert first["state"] == "running" + assert first["events"] == _events("hello", 3)[:2] + assert (first["next"], first["has_more"]) == (2, True) + rest = _trajectory(client, {"task_id": task_id, "after": 2}).json() + assert rest["events"] == _events("hello", 3)[2:] + assert (rest["next"], rest["has_more"]) == (3, False) + assert _trajectory(client, {"task_id": task_id, "after": 3}).json()["events"] == [] + assert _trajectory(client, {"task_id": task_id, "after": 4}).status_code == 400 + assert _trajectory(client, {"task_id": task_id}).status_code == 404 + + gates.release.set() + assert _within(5, lambda: _task_state(client, task_id) == "completed") + final = _trajectory(client, {"task_id": task_id}) + assert final.json() == {"trajectory": _events("hello", 3)} + done = _trajectory(client, {"task_id": task_id, "after": 0}).json() + assert done["state"] == "completed" + assert done["events"] == _events("hello", 3) + + +@pytest.mark.parametrize( + "payload", + [ + {"task_id": "unknown", "after": 0}, + {"context_id": "unknown", "after": 0}, + {"context_id": "unknown"}, + ], +) +def test_unknown_ids_are_not_found(payload: dict[str, Any]) -> None: + with TestClient(_live_agent(_Gates()).create_app()) as client: + assert _trajectory(client, payload).status_code == 404 + + +@pytest.mark.parametrize( + "payload", + [ + {"task_id": "task", "after": True}, + {"task_id": "task", "after": "3"}, + {"task_id": "task", "after": -1}, + {"task_id": "task", "after": 0, "limit": 0}, + ], +) +def test_cursor_fields_are_strict(payload: dict[str, Any]) -> None: + with TestClient(_live_agent(_Gates()).create_app()) as client: + assert _trajectory(client, payload).status_code == 400 + + +@pytest.mark.parametrize(("result", "a2a_state"), [("failed", "failed"), ("raise", "failed")]) +def test_a_failed_task_keeps_its_events(result: str, a2a_state: str) -> None: + with TestClient(_live_agent(_Gates(), result=result).create_app()) as client: + task_id = _send(client, "hello", "message-1") + assert _within(5, lambda: _task_state(client, task_id) == a2a_state) + read = _trajectory(client, {"task_id": task_id, "after": 0}).json() + assert read["state"] == "failed" + assert read["events"] == _events("hello", 3) + assert _trajectory(client, {"task_id": task_id}).json() == { + "trajectory": _events("hello", 3) + } + + +def test_cancel_seals_the_log_canceled_with_its_cleanup_events() -> None: + @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) + class Agent(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.append({"step": "working"}) + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + request.trajectory.append({"step": "cleanup"}) + raise + return TaskResult.text("unreachable") + + def read(client: TestClient, task_id: str) -> dict[str, Any]: + return _trajectory(client, {"task_id": task_id, "after": 0}).json() + + with TestClient(Agent().create_app()) as client: + task_id = _send(client, "hang", "message-1") + assert _within(5, lambda: read(client, task_id)["state"] == "running") + + assert _cancel(client, task_id)["result"]["status"]["state"] == "canceled" + + assert _within(5, lambda: read(client, task_id)["state"] == "canceled") + assert read(client, task_id)["events"] == [{"step": "working"}, {"step": "cleanup"}] + + +def test_a_cancel_the_agent_swallows_still_seals_canceled() -> None: + services = _SdkServices((enable(TRAJECTORY_V1, live=True),), None) + trajectories = services.task_trajectories + log = trajectories.register("task-1", "context-1") + trajectories.start("task-1") + log.append({"step": 1}) + + trajectories.request_cancel("task-1") + trajectories.seal("task-1", TrajectoryState.COMPLETED) + + page = trajectories.task_page("task-1", 0, 10, 1_000) + assert page is not None and page.state is TrajectoryState.CANCELED + + +def test_a_native_trajectory_is_the_final_record(caplog: pytest.LogCaptureFixture) -> None: + @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) + class Agent(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.append({"live": True}) + return ( + TaskResult.builder() + .succeeded() + .add_text("done") + .native_trajectory(format="native/1", payload=[{"native": True}]) + .build() + ) + + with caplog.at_level(logging.WARNING), TestClient(Agent().create_app()) as client: + task_id = _send(client, "hello", "message-1") + assert _within(5, lambda: _task_state(client, task_id) == "completed") + + assert _trajectory(client, {"task_id": task_id}).json() == { + "trajectory": [{"native": True}] + } + assert _trajectory(client, {"task_id": task_id, "after": 0}).json()["events"] == [ + {"live": True} + ] + assert "the native trajectory is its final record" in caplog.text + + +def test_context_reads_span_turns_in_order() -> None: + with TestClient(_live_agent(_Gates(), events=2).create_app()) as client: + first = _send(client, "one", "message-1") + assert _within(5, lambda: _task_state(client, first) == "completed") + second = _send(client, "two", "message-2") + assert _within(5, lambda: _task_state(client, second) == "completed") + + everything = _events("one", 2) + _events("two", 2) + read = _trajectory(client, {"context_id": "context-1", "after": 0}).json() + assert read == { + "context_id": "context-1", + "task_id": second, + "state": "completed", + "format": "test-events/1", + "events": everything, + "next": 4, + "has_more": False, + } + tail = _trajectory(client, {"context_id": "context-1", "after": 1, "limit": 2}).json() + assert tail["events"] == everything[1:3] + assert (tail["next"], tail["has_more"]) == (3, True) + assert _trajectory(client, {"context_id": "context-1"}).json() == { + "trajectory": everything + } + + +def test_an_agent_context_handler_wins_over_the_framework() -> None: + @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) + class Agent(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.append({"live": True}) + return TaskResult.text("done") + + @extension(TRAJECTORY_V1.get.context) + async def get_context(self, request: ContextTrajectoryRequest) -> dict[str, Any]: + return {"trajectory": "the agent's own"} + + with TestClient(Agent().create_app()) as client: + task_id = _send(client, "hello", "message-1") + assert _within(5, lambda: _task_state(client, task_id) == "completed") + + assert _trajectory(client, {"context_id": "context-1"}).json() == { + "trajectory": "the agent's own" + } + read = _trajectory(client, {"context_id": "context-1", "after": 0}).json() + assert read["events"] == [{"live": True}] + + +def test_a_streaming_run_appends_the_same_way() -> None: + @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) + class Agent(AgentEnvAgent): + async def run(self, request: TaskRequest) -> AsyncIterator[Any]: + request.trajectory.append({"step": 1}) + yield TaskProgress.text("working") + request.trajectory.append({"step": 2}) + yield TaskResult.text("done") + + with TestClient(Agent().create_app()) as client: + task_id = _send(client, "hello", "message-1") + assert _within(5, lambda: _task_state(client, task_id) == "completed") + read = _trajectory(client, {"task_id": task_id, "after": 0}).json() + assert read["state"] == "completed" + assert read["events"] == [{"step": 1}, {"step": 2}] + + +@pytest.mark.asyncio +async def test_default_handlers_delegate_the_cursor_read() -> None: + services = _SdkServices((enable(TRAJECTORY_V1, live=True),), None) + log = services.task_trajectories.register("task-1", "context-1") + services.task_trajectories.start("task-1") + log.append({"step": 1}) + + result = await DefaultExtensionHandlers(services).call( + TRAJECTORY_V1.get, TaskEventsTrajectoryRequest(task_id="task-1", after=0) + ) + + assert result.events == [{"step": 1}] + assert result.state is TrajectoryState.RUNNING + + +def test_an_appended_event_never_changes() -> None: + log = TrajectoryLog() + event = {"nested": {"value": 1}} + log.append(event) + event["nested"]["value"] = 2 + + assert log._events == [b'{"nested":{"value":1}}'] + for bad in ({"value": float("nan")}, {"value": {1, 2}}): + with pytest.raises((TypeError, ValueError)): + log.append(bad) + with pytest.raises(TypeError, match="JSON object"): + log.append(["not", "an", "object"]) # type: ignore[arg-type] + assert len(log) == 1 + with pytest.raises(ValueError, match="non-empty"): + log.set_format(" ") + + +def test_a_sealed_log_refuses_more_events() -> None: + trajectories = TaskTrajectories() + log = trajectories.register("task-1", "context-1") + log.append({"step": 1}) + trajectories.seal("task-1", TrajectoryState.COMPLETED) + + with pytest.raises(RuntimeError, match="can no longer change"): + log.append({"step": 2}) + with pytest.raises(RuntimeError, match="can no longer change"): + log.set_format("late/1") + assert trajectories.final("task-1") == b'[{"step":1}]' + + +def test_reads_never_return_a_partial_event() -> None: + trajectories = TaskTrajectories() + log = trajectories.register("task-1", "context-1") + trajectories.start("task-1") + seen: list[bytes] = [] + for index in range(50): + log.append({"index": index, "padding": "x" * index}) + page = trajectories.task_page("task-1", len(seen), 7, 1_000) + assert page is not None + seen.extend(page.events) + while len(seen) < 50: + page = trajectories.task_page("task-1", len(seen), 7, 1_000) + assert page is not None + seen.extend(page.events) + + assert [json.loads(item) for item in seen] == [ + {"index": index, "padding": "x" * index} for index in range(50) + ] + + +def test_the_byte_budget_still_returns_one_event() -> None: + trajectories = TaskTrajectories() + log = trajectories.register("task-1", "context-1") + trajectories.start("task-1") + log.append({"big": "x" * 100}) + log.append({"big": "y" * 100}) + + page = trajectories.task_page("task-1", 0, 10, 10) + + assert page is not None and len(page.events) == 1 + + +def test_a_task_that_starts_later_moves_without_shifting_positions() -> None: + trajectories = TaskTrajectories() + queued = trajectories.register("queued", "context-1") + running = trajectories.register("running", "context-1") + trajectories.start("running") + running.append({"from": "running"}) + trajectories.start("queued") + queued.append({"from": "queued"}) + + page = trajectories.context_page("context-1", 0, 10, 1_000) + + assert page is not None + assert [json.loads(item) for item in page.events] == [ + {"from": "running"}, + {"from": "queued"}, + ] + assert page.task_id == "running" + + +def test_the_byte_cap_evicts_completed_contexts_but_never_running_ones() -> None: + trajectories = TaskTrajectories(max_bytes=50) + running = trajectories.register("running", "context-running") + trajectories.start("running") + running.append({"padding": "x" * 60}) + for number in (1, 2): + trajectories.register(f"done-{number}", f"context-{number}") + trajectories.seal( + f"done-{number}", TrajectoryState.COMPLETED, native=b"[" + b"1" * 30 + b"]" + ) + + assert trajectories.task_page("running", 0, 10, 1_000) is not None + assert trajectories.final("done-1") is None + assert trajectories.final("done-2") is None + + +def test_a_large_running_trajectory_warns_once( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(tasks_v1, "_LARGE_LIVE_TRAJECTORY_BYTES", 10) + log = TaskTrajectories().register("task-1", "context-1") + + with caplog.at_level(logging.WARNING): + for _ in range(3): + log.append({"padding": "x" * 20}) + + assert caplog.text.count("trajectory of running task task-1 exceeds") == 1 From 728dc64e6045637f96e2c4b9cab4f96fa50708d5 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 7 Oct 2026 15:32:01 -0700 Subject: [PATCH 02/12] fix(protocol): bound cursor limit and keep a just-ended context through eviction - `limit` on the cursor reads is validated as 1 to 1,000 (HTTP 400 outside), instead of being quietly lowered to the cap. The cap moves to `extensions.py` as `MAX_TRAJECTORY_READ_EVENTS`, next to the request models it bounds. - The trajectory store evicts by count only (1,024 contexts, least recently used first), so finished trajectories are kept at least as long as the previous 1,024-task cache kept them; the 256 MiB byte cap is removed. - Eviction never drops the context whose task has just ended, so its final record is still there when the client fetches it. Previously a context over the byte cap could be evicted at the moment its task was sealed. FD-3719 Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 13 +++---- .../a2a_agent/_trajectory_log.py | 36 ++++++++----------- .../agentenv_protocol/a2a_agent/extensions.py | 7 ++-- .../agentenv_protocol/a2a_agent/framework.py | 11 ++---- .../tests/test_a2a_trajectory_log.py | 28 +++++++++++---- 5 files changed, 50 insertions(+), 45 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 37beb205..e4a07358 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -347,8 +347,9 @@ Both answer `{context_id, task_id, state, format, events, next, has_more}` - `events` are the events from position `after`, zero-based. `next` is `after` plus the number returned, the position to read from next; `has_more` says more events are already readable past this page. A page holds at most - 1,000 events and about 4 MiB, but always at least one event when one exists; - `limit` lowers the event count. `after` past the last event is HTTP 400. + `limit` events (1 to 1,000, default 1,000) and about 4 MiB, but always at + least one event when one exists. A `limit` outside 1 to 1,000, or `after` past + the last event, is HTTP 400. - `state` is `pending` before the first event, `running` while the task runs, then `completed`, `failed` or `canceled`. A terminal state is reported only once the log can no longer change, so a reader that has reached `has_more: @@ -369,10 +370,10 @@ overrides the whole `get` operation must accept the cursor requests too. Clients detect live support from the card, with `card_request_accepts(get_request, {"task_id", "after"})`. -The logs are kept in memory, up to 256 MiB and 1,024 contexts. Past that, whole -contexts are dropped, least recently used first; a context with a task still -pending or running is never dropped. A running task's log is logged as a warning -once it passes 100 MiB. +The logs are kept in memory for up to 1,024 contexts. Past that, whole contexts +are dropped, least recently used first; a context with a task still pending or +running is never dropped, nor is the one whose task has just ended. A running +task's log is logged as a warning once it passes 100 MiB. `enable(..., description="...")` is reserved for declarations carrying configuration or metadata. It preserves the agent-specific extension prose diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py index 6905eaa4..b5121dc4 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py @@ -11,9 +11,7 @@ logger = logging.getLogger(__name__) -MAX_CACHED_TRAJECTORY_BYTES = 256 * 1024 * 1024 MAX_CACHED_CONTEXTS = 1_024 -MAX_TRAJECTORY_READ_EVENTS = 1_000 MAX_TRAJECTORY_READ_BYTES = 4 * 1024 * 1024 @@ -32,10 +30,6 @@ def state(self) -> TrajectoryState: return self.final_state return TrajectoryState.RUNNING if len(self.log) else TrajectoryState.PENDING - @property - def size_bytes(self) -> int: - return self.log._size_bytes + len(self.native or b"") - @dataclass(frozen=True, slots=True) class TrajectoryPage: @@ -50,17 +44,13 @@ class TrajectoryPage: class TaskTrajectories: """Each task's log and final record, grouped by context in execution order. - A whole context is the unit of eviction, least recently used first, so positions in a - context read never shift; a context with a task pending or running is never evicted.""" + Past ``max_contexts``, whole contexts are evicted, least recently used first, so + positions in a context read never shift. A context with a task pending or running is + never evicted, nor is the one whose task just ended, so its final record is readable.""" - def __init__( - self, - max_bytes: int = MAX_CACHED_TRAJECTORY_BYTES, - max_contexts: int = MAX_CACHED_CONTEXTS, - ) -> None: - if max_bytes < 1 or max_contexts < 1: - raise ValueError("trajectory cache bounds must be positive") - self._max_bytes = max_bytes + def __init__(self, max_contexts: int = MAX_CACHED_CONTEXTS) -> None: + if max_contexts < 1: + raise ValueError("max_contexts must be positive") self._max_contexts = max_contexts self._tasks: dict[str, _TaskRecord] = {} self._contexts: OrderedDict[str, list[str]] = OrderedDict() @@ -113,7 +103,8 @@ def seal( record.native = native record.log._sealed = True record.final_state = TrajectoryState.CANCELED if record.cancel_requested else state - self._evict() + self._touch(record.context_id) + self._evict(keep=record.context_id) def final(self, task_id: str) -> bytes | None: """The ended task's final trajectory as a JSON array, or None if there is none yet.""" @@ -211,17 +202,18 @@ def _touch(self, context_id: str) -> None: if context_id in self._contexts: self._contexts.move_to_end(context_id) - def _evict(self) -> None: - total = sum(record.size_bytes for record in self._tasks.values()) + def _evict(self, *, keep: str) -> None: for context_id in tuple(self._contexts): - if total <= self._max_bytes and len(self._contexts) <= self._max_contexts: + if len(self._contexts) <= self._max_contexts: return task_ids = self._contexts[context_id] - if any(self._tasks[task_id].final_state is None for task_id in task_ids): + if context_id == keep or any( + self._tasks[task_id].final_state is None for task_id in task_ids + ): continue del self._contexts[context_id] for task_id in task_ids: - total -= self._tasks.pop(task_id).size_bytes + del self._tasks[task_id] def _page( diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py index 0a049002..a0a7db36 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py @@ -144,16 +144,19 @@ class TrajectoryState(str, Enum): CANCELED = "canceled" +MAX_TRAJECTORY_READ_EVENTS = 1_000 + + class TaskEventsTrajectoryRequest(ExtensionRequest): task_id: str after: StrictInt = Field(ge=0) - limit: StrictInt | None = Field(default=None, ge=1) + limit: StrictInt | None = Field(default=None, ge=1, le=MAX_TRAJECTORY_READ_EVENTS) class ContextEventsTrajectoryRequest(ExtensionRequest): context_id: str after: StrictInt = Field(ge=0) - limit: StrictInt | None = Field(default=None, ge=1) + limit: StrictInt | None = Field(default=None, ge=1, le=MAX_TRAJECTORY_READ_EVENTS) class TrajectoryEventsResponse(ExtensionResponse): diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py index 84fd8a07..25fcc74d 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/framework.py @@ -28,14 +28,11 @@ from starlette.routing import Route from starlette.types import ASGIApp, Receive, Scope, Send -from ._trajectory_log import ( - MAX_TRAJECTORY_READ_BYTES, - MAX_TRAJECTORY_READ_EVENTS, - TaskTrajectories, -) +from ._trajectory_log import MAX_TRAJECTORY_READ_BYTES, TaskTrajectories from ._triggers import TriggerEngine, TriggerError from .extensions import ( AGENT_CONFIG_V1, + MAX_TRAJECTORY_READ_EVENTS, MCP_CONFIG_V1, SKILL_CONFIG_V1, TRAJECTORY_V1, @@ -770,9 +767,7 @@ async def trajectory_get( def _trajectory_events( self, request: TaskEventsTrajectoryRequest | ContextEventsTrajectoryRequest ) -> TrajectoryEventsResponse: - max_events = min( - request.limit or MAX_TRAJECTORY_READ_EVENTS, MAX_TRAJECTORY_READ_EVENTS - ) + max_events = request.limit or MAX_TRAJECTORY_READ_EVENTS if isinstance(request, TaskEventsTrajectoryRequest): page = self.task_trajectories.task_page( request.task_id, request.after, max_events, MAX_TRAJECTORY_READ_BYTES diff --git a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py index a948e69f..8f44e5f9 100644 --- a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py +++ b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py @@ -227,6 +227,8 @@ def test_unknown_ids_are_not_found(payload: dict[str, Any]) -> None: {"task_id": "task", "after": "3"}, {"task_id": "task", "after": -1}, {"task_id": "task", "after": 0, "limit": 0}, + {"task_id": "task", "after": 0, "limit": 1_001}, + {"context_id": "context", "after": 0, "limit": 1_001}, ], ) def test_cursor_fields_are_strict(payload: dict[str, Any]) -> None: @@ -473,20 +475,32 @@ def test_a_task_that_starts_later_moves_without_shifting_positions() -> None: assert page.task_id == "running" -def test_the_byte_cap_evicts_completed_contexts_but_never_running_ones() -> None: - trajectories = TaskTrajectories(max_bytes=50) +def test_eviction_drops_completed_contexts_but_never_running_ones() -> None: + trajectories = TaskTrajectories(max_contexts=2) running = trajectories.register("running", "context-running") trajectories.start("running") - running.append({"padding": "x" * 60}) - for number in (1, 2): + running.append({"step": 1}) + for number in (1, 2, 3): trajectories.register(f"done-{number}", f"context-{number}") - trajectories.seal( - f"done-{number}", TrajectoryState.COMPLETED, native=b"[" + b"1" * 30 + b"]" - ) + trajectories.seal(f"done-{number}", TrajectoryState.COMPLETED, native=b"[1]") assert trajectories.task_page("running", 0, 10, 1_000) is not None assert trajectories.final("done-1") is None assert trajectories.final("done-2") is None + assert trajectories.final("done-3") == b"[1]" + + +def test_the_context_whose_task_just_ended_is_kept_for_its_final_read() -> None: + trajectories = TaskTrajectories(max_contexts=1) + log = trajectories.register("ending", "context-ending") + trajectories.start("ending") + log.append({"step": 1}) + for number in (1, 2): + trajectories.register(f"queued-{number}", f"context-queued-{number}") + + trajectories.seal("ending", TrajectoryState.COMPLETED) + + assert trajectories.final("ending") == b'[{"step":1}]' def test_a_large_running_trajectory_warns_once( From f4e02f69749eb7db9aeedb9d2e738f960ea0b6dd Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 7 Oct 2026 17:16:01 -0700 Subject: [PATCH 03/12] feat(protocol): require a trajectory format before the first append `TrajectoryLog.append` raises until `set_format` has named the events' format, so every live trajectory carries one for readers to parse it by. A context read reports the most recent format any of its tasks set, so a turn that has not started yet no longer hides the earlier turns' format. FD-3719 Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 5 +- .../a2a_agent/_trajectory_log.py | 5 +- .../agentenv_protocol/a2a_agent/tasks/v1.py | 4 +- .../tests/test_a2a_trajectory_log.py | 58 +++++++++++++++---- 4 files changed, 58 insertions(+), 14 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index e4a07358..598f0da4 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -326,7 +326,7 @@ class MyAgent(AgentEnvAgent): never changes after it is appended and a read never returns part of one. It does no I/O; call it from one thread at a time. After the task ends the log refuses further events. `set_format` names the events' versioned format, which readers -use to parse them. +use to parse them; call it before the first `append`, which otherwise raises. The framework keeps each task's log and serves it while the task runs. When the result carries no `native_trajectory`, the log is the task's final trajectory: @@ -357,7 +357,8 @@ Both answer `{context_id, task_id, state, format, events, next, has_more}` before the task's terminal A2A status, so the completed-task reads succeed as soon as a client sees the task end. - A context read reports the context's current task, the running one or else the - latest, with its `format`; its `state` is `running` while any task runs. + latest, and the most recent `format` any of its tasks set; its `state` is + `running` while any task runs. - A task or context the agent has not seen, or no longer keeps, is HTTP 404. A task is known from the moment `message/send` returns its id, so a 404 for an id that was readable means it was evicted, and any cursor held for it, or for diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py index b5121dc4..0240712f 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py @@ -186,7 +186,10 @@ def context_page( context_id=context_id, task_id=task_ids[current], state=state, - format=records[current].log.format, + format=next( + (record.log.format for record in reversed(records) if record.log.format), + None, + ), events=page, total=sum(lengths), ) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py index 657121da..2503f798 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py @@ -159,10 +159,12 @@ def set_format(self, format: str) -> None: self._format = format def append(self, event: Mapping[str, Any]) -> None: - """Record one complete event, a JSON object.""" + """Record one complete event, a JSON object, after ``set_format`` has named its format.""" if not isinstance(event, Mapping): raise TypeError("a trajectory event must be a JSON object") self._ensure_open() + if self._format is None: + raise RuntimeError("call set_format(...) before appending trajectory events") encoded = json.dumps(dict(event), separators=(",", ":"), allow_nan=False).encode() self._events.append(encoded) self._size_bytes += len(encoded) diff --git a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py index 8f44e5f9..8cd1d4ad 100644 --- a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py +++ b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py @@ -128,6 +128,12 @@ def _within(seconds: float, condition) -> bool: return True +def _open(trajectories: TaskTrajectories, task_id: str, context_id: str) -> TrajectoryLog: + log = trajectories.register(task_id, context_id) + log.set_format("test-events/1") + return log + + def _events(text: str, count: int) -> list[dict[str, Any]]: return [{"text": text, "index": index} for index in range(count)] @@ -253,6 +259,7 @@ def test_cancel_seals_the_log_canceled_with_its_cleanup_events() -> None: @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) class Agent(AgentEnvAgent): async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.set_format("test-events/1") request.trajectory.append({"step": "working"}) try: await asyncio.Event().wait() @@ -277,7 +284,7 @@ def read(client: TestClient, task_id: str) -> dict[str, Any]: def test_a_cancel_the_agent_swallows_still_seals_canceled() -> None: services = _SdkServices((enable(TRAJECTORY_V1, live=True),), None) trajectories = services.task_trajectories - log = trajectories.register("task-1", "context-1") + log = _open(trajectories, "task-1", "context-1") trajectories.start("task-1") log.append({"step": 1}) @@ -292,6 +299,7 @@ def test_a_native_trajectory_is_the_final_record(caplog: pytest.LogCaptureFixtur @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) class Agent(AgentEnvAgent): async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.set_format("test-events/1") request.trajectory.append({"live": True}) return ( TaskResult.builder() @@ -344,6 +352,7 @@ def test_an_agent_context_handler_wins_over_the_framework() -> None: @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) class Agent(AgentEnvAgent): async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.set_format("test-events/1") request.trajectory.append({"live": True}) return TaskResult.text("done") @@ -366,6 +375,7 @@ def test_a_streaming_run_appends_the_same_way() -> None: @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) class Agent(AgentEnvAgent): async def run(self, request: TaskRequest) -> AsyncIterator[Any]: + request.trajectory.set_format("test-events/1") request.trajectory.append({"step": 1}) yield TaskProgress.text("working") request.trajectory.append({"step": 2}) @@ -382,7 +392,7 @@ async def run(self, request: TaskRequest) -> AsyncIterator[Any]: @pytest.mark.asyncio async def test_default_handlers_delegate_the_cursor_read() -> None: services = _SdkServices((enable(TRAJECTORY_V1, live=True),), None) - log = services.task_trajectories.register("task-1", "context-1") + log = _open(services.task_trajectories, "task-1", "context-1") services.task_trajectories.start("task-1") log.append({"step": 1}) @@ -396,6 +406,7 @@ async def test_default_handlers_delegate_the_cursor_read() -> None: def test_an_appended_event_never_changes() -> None: log = TrajectoryLog() + log.set_format("test-events/1") event = {"nested": {"value": 1}} log.append(event) event["nested"]["value"] = 2 @@ -413,7 +424,7 @@ def test_an_appended_event_never_changes() -> None: def test_a_sealed_log_refuses_more_events() -> None: trajectories = TaskTrajectories() - log = trajectories.register("task-1", "context-1") + log = _open(trajectories, "task-1", "context-1") log.append({"step": 1}) trajectories.seal("task-1", TrajectoryState.COMPLETED) @@ -426,7 +437,7 @@ def test_a_sealed_log_refuses_more_events() -> None: def test_reads_never_return_a_partial_event() -> None: trajectories = TaskTrajectories() - log = trajectories.register("task-1", "context-1") + log = _open(trajectories, "task-1", "context-1") trajectories.start("task-1") seen: list[bytes] = [] for index in range(50): @@ -446,7 +457,7 @@ def test_reads_never_return_a_partial_event() -> None: def test_the_byte_budget_still_returns_one_event() -> None: trajectories = TaskTrajectories() - log = trajectories.register("task-1", "context-1") + log = _open(trajectories, "task-1", "context-1") trajectories.start("task-1") log.append({"big": "x" * 100}) log.append({"big": "y" * 100}) @@ -458,8 +469,8 @@ def test_the_byte_budget_still_returns_one_event() -> None: def test_a_task_that_starts_later_moves_without_shifting_positions() -> None: trajectories = TaskTrajectories() - queued = trajectories.register("queued", "context-1") - running = trajectories.register("running", "context-1") + queued = _open(trajectories, "queued", "context-1") + running = _open(trajectories, "running", "context-1") trajectories.start("running") running.append({"from": "running"}) trajectories.start("queued") @@ -477,7 +488,7 @@ def test_a_task_that_starts_later_moves_without_shifting_positions() -> None: def test_eviction_drops_completed_contexts_but_never_running_ones() -> None: trajectories = TaskTrajectories(max_contexts=2) - running = trajectories.register("running", "context-running") + running = _open(trajectories, "running", "context-running") trajectories.start("running") running.append({"step": 1}) for number in (1, 2, 3): @@ -492,7 +503,7 @@ def test_eviction_drops_completed_contexts_but_never_running_ones() -> None: def test_the_context_whose_task_just_ended_is_kept_for_its_final_read() -> None: trajectories = TaskTrajectories(max_contexts=1) - log = trajectories.register("ending", "context-ending") + log = _open(trajectories, "ending", "context-ending") trajectories.start("ending") log.append({"step": 1}) for number in (1, 2): @@ -507,10 +518,37 @@ def test_a_large_running_trajectory_warns_once( caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(tasks_v1, "_LARGE_LIVE_TRAJECTORY_BYTES", 10) - log = TaskTrajectories().register("task-1", "context-1") + log = _open(TaskTrajectories(), "task-1", "context-1") with caplog.at_level(logging.WARNING): for _ in range(3): log.append({"padding": "x" * 20}) assert caplog.text.count("trajectory of running task task-1 exceeds") == 1 + + +def test_appending_before_set_format_is_refused() -> None: + log = TrajectoryLog() + + with pytest.raises(RuntimeError, match="set_format"): + log.append({"step": 1}) + assert len(log) == 0 + + +def test_a_context_read_keeps_the_format_while_a_new_turn_is_pending() -> None: + trajectories = TaskTrajectories() + first = _open(trajectories, "first", "context-1") + trajectories.start("first") + first.append({"step": 1}) + trajectories.seal("first", TrajectoryState.COMPLETED) + trajectories.register("second", "context-1") + trajectories.start("second") + + page = trajectories.context_page("context-1", 0, 10, 1_000) + + assert page is not None + assert (page.task_id, page.state, page.format) == ( + "second", + TrajectoryState.PENDING, + "test-events/1", + ) From f14f861d8ea2a8946f64f05b3b3cb3eba6d2c352 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 7 Oct 2026 17:16:01 -0700 Subject: [PATCH 04/12] feat(prompt_agent): follow a turn's live trajectory and store it in chunks `prompt_agent` gains `live_trajectory` (default off; also settable per run through `step_params`). When on and the agent's card offers the cursor read, each turn's trajectory is followed while the agent works and stored under the turn's trajectory prefix: /live/-.jsonl per page /live/meta.json {"format"} /live/end.json {"state", "next"} - `agent_env.a2a_agent.trajectory_follower`: `live_trajectory_endpoint` reads the card; `follow_trajectory` reads one task from its first event, hands each page to a sink, retries transport errors and 5xx, stops on 404, and returns once the task has ended and every event is read. - `snapshot_utils.live_trajectory`: the chunk sink (write-once objects, so a repeated write is a no-op) and `following_live_trajectory`, which drains after a turn that ended normally (at most 60 s) and stops at once after one that failed. A follower's failure never fails the turn. - `send_and_wait` takes `on_sent`, called with the peer's task id. - With the flag off, an agent without the cursor read, or a prefix outside the configured store, the step runs exactly as before; `to_dict` records `live_trajectory` only when it is on. FD-3719 Co-Authored-By: Claude Opus 5.5 --- src/agent_env/a2a_agent/object_transfer.py | 8 +- .../a2a_agent/trajectory_follower.py | 123 +++++++++++ .../snapshot_utils/agent_state_capture.py | 19 +- .../snapshot_utils/live_trajectory.py | 109 +++++++++ .../task_step/task_steps/prompt_agent.py | 49 ++++- .../a2a_agent/test_trajectory_follower.py | 206 ++++++++++++++++++ .../test_prompt_agent_live_trajectory.py | 174 +++++++++++++++ 7 files changed, 674 insertions(+), 14 deletions(-) create mode 100644 src/agent_env/a2a_agent/trajectory_follower.py create mode 100644 src/agent_env/task_step/snapshot_utils/live_trajectory.py create mode 100644 tst/unit/a2a_agent/test_trajectory_follower.py create mode 100644 tst/unit/task_step/test_prompt_agent_live_trajectory.py diff --git a/src/agent_env/a2a_agent/object_transfer.py b/src/agent_env/a2a_agent/object_transfer.py index 034105ed..8fff9e46 100644 --- a/src/agent_env/a2a_agent/object_transfer.py +++ b/src/agent_env/a2a_agent/object_transfer.py @@ -359,14 +359,16 @@ async def send_and_wait( timeout_seconds: int, poll_interval_seconds: int, before_send: Callable[[], None] | None = None, + on_sent: Callable[[str], None] | None = None, shareable: Collection[str] | None = None, ) -> tuple[str, dict]: """Send ``parts`` to the A2A peer at ``a2a_url`` and wait for its task to end: the task's id and its final state. A peer that is an ``agent`` agent-env deployed is sent each file part a configured store owns, of those ``shareable`` names when given, as an HTTPS URL it can read (``readable_parts``); any other peer, such as a human's hub, reads the store itself and is sent the parts as they are. - ``before_send`` runs once the parts are ready, just before the message goes out. An ``agent``'s sandbox is - watched while it works, so one that dies is given up on (``poll_a2a_task``).""" + ``before_send`` runs once the parts are ready, just before the message goes out, and ``on_sent`` with the + peer's task id once it is accepted. An ``agent``'s sandbox is watched while it works, so one that dies is + given up on (``poll_a2a_task``).""" sending = ( readable_parts( parts, a2a_url=a2a_url, card=agent.a2a_card, sandbox_type=agent.sandbox_type, @@ -379,6 +381,8 @@ async def send_and_wait( if before_send is not None: before_send() task_id, _ = await protocol.send_a2a_message(a2a_url, sent, message_id, context_id, timeout_seconds) + if on_sent is not None: + on_sent(task_id) return task_id, await protocol.poll_a2a_task( a2a_url, task_id, timeout_seconds, poll_interval_seconds, sandbox_id=agent.sandbox_id if agent is not None else None, diff --git a/src/agent_env/a2a_agent/trajectory_follower.py b/src/agent_env/a2a_agent/trajectory_follower.py new file mode 100644 index 00000000..9bcaed79 --- /dev/null +++ b/src/agent_env/a2a_agent/trajectory_follower.py @@ -0,0 +1,123 @@ +"""Follow one A2A task's trajectory while it runs, through the trajectory extension's cursor read.""" +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Awaitable, Callable +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Any + +import httpx +from agentenv_protocol.a2a_agent import ( + TaskEventsTrajectoryRequest, + TrajectoryEventsResponse, + TrajectoryState, + card_request_accepts, + request_fields, +) + +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.a2a_agent.object_transfer import parse_response +from agent_env.a2a_agent.protocol import raise_for_extension_status + +logger = logging.getLogger(__name__) + +_CURSOR_FIELDS = frozenset(request_fields(TaskEventsTrajectoryRequest).required) +_OPEN_STATES = frozenset({TrajectoryState.PENDING, TrajectoryState.RUNNING}) +_READ_TIMEOUT_SECONDS = 30 +_MAX_BACKOFF_SECONDS = 60 +# Once stopped, the task has ended and its log is sealed, so the next read is normally the last. +_STOPPED_INTERVAL_SECONDS = 0.2 + + +@dataclass(frozen=True) +class TrajectoryBatch: + """Events ``after`` to ``next`` of a task's trajectory, as one read returned them; ``last`` + when the task has ended and no event remains unread.""" + + after: int + next: int + events: list[Any] + format: str | None + state: TrajectoryState + last: bool + + +TrajectorySink = Callable[[TrajectoryBatch], Awaitable[None]] + + +def live_trajectory_endpoint(card: dict, a2a_url: str) -> str | None: + """The URL of the agent's live trajectory read, or None when its card does not offer one.""" + extension = A2AAgent.find_extension(card, A2AAgent.EXT_TRAJECTORY) + method, path = A2AAgent.operation(extension, "get") + if not method or path is None: + return None + if not card_request_accepts(method.get("request") or {}, _CURSOR_FIELDS): + return None + return a2a_url + path + + +async def follow_trajectory( + endpoint: str, + task_id: str, + sink: TrajectorySink, + *, + poll_interval_seconds: float, + stop: asyncio.Event, + client: httpx.AsyncClient | None = None, +) -> None: + """Read task ``task_id``'s trajectory from its first event, handing ``sink`` each page that has + events and the last one, until the task has ended and every event is read. + + Between reads it waits ``poll_interval_seconds``, or less once ``stop`` is set. Transport errors + and 5xx answers are retried with backoff; a task the agent no longer knows (404) ends following.""" + after = 0 + failures = 0 + async with nullcontext(client) if client is not None else httpx.AsyncClient() as http: + while True: + try: + response = await http.post( + endpoint, json={"task_id": task_id, "after": after}, timeout=_READ_TIMEOUT_SECONDS + ) + if response.status_code == 404: + logger.warning("Agent no longer knows task %s; stopped following its trajectory", task_id) + return + raise_for_extension_status(response, operation="trajectory get", include_body=True) + except httpx.HTTPStatusError as exc: + if exc.response.status_code < 500: + raise + failures += 1 + logger.warning("Live trajectory read failed (consec=%d, will retry): %s", failures, exc) + except httpx.TransportError as exc: + failures += 1 + logger.warning("Live trajectory read failed (consec=%d, will retry): %r", failures, exc) + else: + failures = 0 + page = parse_response(TrajectoryEventsResponse, response.json(), operation="trajectory get") + last = page.state not in _OPEN_STATES and not page.has_more + if page.events or last: + await sink( + TrajectoryBatch( + after=after, + next=page.next, + events=page.events, + format=page.format, + state=page.state, + last=last, + ) + ) + after = page.next + if last: + return + if page.has_more: + continue + if stop.is_set(): + await asyncio.sleep(_STOPPED_INTERVAL_SECONDS) + continue + try: + await asyncio.wait_for( + stop.wait(), timeout=min(poll_interval_seconds + failures * 5, _MAX_BACKOFF_SECONDS) + ) + except TimeoutError: + pass diff --git a/src/agent_env/task_step/snapshot_utils/agent_state_capture.py b/src/agent_env/task_step/snapshot_utils/agent_state_capture.py index 0ab32b82..3630b945 100644 --- a/src/agent_env/task_step/snapshot_utils/agent_state_capture.py +++ b/src/agent_env/task_step/snapshot_utils/agent_state_capture.py @@ -223,12 +223,7 @@ def upload_trajectory( """ store = get_config().get_object_store() - # Via the store, not urlparse().path: a prefix naming a different bucket - # must fail loudly, not silently write to the configured one. - prefix_key = store.get_object_key(trajectory_output_prefix) - if prefix_key and not prefix_key.endswith("/"): - prefix_key += "/" - key = f"{prefix_key}trajectory-{name or uuid.uuid4().hex[:12]}.json" + key = f"{trajectory_prefix_key(store, trajectory_output_prefix)}trajectory-{name or uuid.uuid4().hex[:12]}.json" body = json.dumps(trajectory, indent=2, default=str).encode() return store.put(key, body, content_type="application/json") @@ -248,8 +243,16 @@ def trajectory_object_url( trajectory_output_prefix: str, *, store, name: Optional[str] = None ) -> str: """Return the durable URL used for one agent-written trajectory.""" + key = f"{trajectory_prefix_key(store, trajectory_output_prefix)}trajectory-{name or uuid.uuid4().hex[:12]}.json" + return store.object_url(key) + + +def trajectory_prefix_key(store, trajectory_output_prefix: str) -> str: + """The store key of a trajectory prefix, ending in ``/``. + + Via the store, not urlparse().path: a prefix naming a different bucket must fail loudly + (``ValueError``), not silently write to the configured one.""" prefix_key = store.get_object_key(trajectory_output_prefix) if prefix_key and not prefix_key.endswith("/"): prefix_key += "/" - key = f"{prefix_key}trajectory-{name or uuid.uuid4().hex[:12]}.json" - return store.object_url(key) + return prefix_key diff --git a/src/agent_env/task_step/snapshot_utils/live_trajectory.py b/src/agent_env/task_step/snapshot_utils/live_trajectory.py new file mode 100644 index 00000000..777e4088 --- /dev/null +++ b/src/agent_env/task_step/snapshot_utils/live_trajectory.py @@ -0,0 +1,109 @@ +"""A turn's live trajectory, stored as it is read: ``/live/-.jsonl`` per +page, ``meta.json`` naming the events' format, and ``end.json`` once the turn has ended. + +Written from the worker with its own credentials, like the final trajectory. Every object is +written once: one already there holds the same bytes, so a repeated write is a no-op.""" +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager + +from agent_env.a2a_agent.trajectory_follower import TrajectoryBatch, follow_trajectory +from agent_env.store.base import ObjectAlreadyExistsError +from agent_env.store.object_store import ObjectStore + +logger = logging.getLogger(__name__) + +_WRITE_ATTEMPTS = 3 +# The task has ended when the turn's wait returns, so draining its last events is normally one read. +_DRAIN_SECONDS = 60 + + +class LiveTrajectoryChunks: + """The ``follow_trajectory`` sink that stores one turn's pages.""" + + def __init__(self, store: ObjectStore, prefix_key: str, turn_id: str) -> None: + self._store = store + self._directory = f"{prefix_key}{turn_id}/live/" + self._format_written = False + + async def __call__(self, batch: TrajectoryBatch) -> None: + if batch.events: + body = b"".join( + json.dumps(event, separators=(",", ":")).encode() + b"\n" for event in batch.events + ) + await self._put(f"{batch.after:08d}-{batch.next:08d}.jsonl", body, "application/x-ndjson") + if batch.format and not self._format_written: + await self._put("meta.json", json.dumps({"format": batch.format}).encode(), "application/json") + self._format_written = True + if batch.last: + end = {"state": batch.state.value, "next": batch.next} + await self._put("end.json", json.dumps(end).encode(), "application/json") + + async def _put(self, name: str, body: bytes, content_type: str) -> None: + key = self._directory + name + for attempt in range(1, _WRITE_ATTEMPTS + 1): + try: + await asyncio.to_thread(self._store.put, key, body, content_type=content_type) + return + except ObjectAlreadyExistsError: + return + except Exception as exc: + if attempt == _WRITE_ATTEMPTS: + raise + logger.warning("Live trajectory write of %s failed (attempt %d, will retry): %s", key, attempt, exc) + await asyncio.sleep(attempt) + + +@asynccontextmanager +async def following_live_trajectory( + endpoint: str, + store: ObjectStore, + prefix_key: str, + turn_id: str, + *, + poll_interval_seconds: float, +) -> AsyncIterator[Callable[[str], None]]: + """Yield the callback that starts following the agent's task once the turn's message is sent. + + When the body returns, the follower drains what remains, for at most ``_DRAIN_SECONDS``; when it + raises, the follower is stopped at once, since the agent may be gone. A follower's own failure + is logged and never fails the turn.""" + stop = asyncio.Event() + followers: list[asyncio.Task] = [] + + def follow(task_id: str) -> None: + sink = LiveTrajectoryChunks(store, prefix_key, turn_id) + followers.append( + asyncio.create_task( + follow_trajectory( + endpoint, task_id, sink, poll_interval_seconds=poll_interval_seconds, stop=stop + ) + ) + ) + + try: + yield follow + stop.set() + if followers: + _, draining = await asyncio.wait(followers, timeout=_DRAIN_SECONDS) + if draining: + logger.warning( + "Live trajectory of turn %s did not finish draining in %ds (continuing)", + turn_id, + _DRAIN_SECONDS, + ) + finally: + stop.set() + for follower in followers: + follower.cancel() + if followers: + await asyncio.wait(followers) + for follower in followers: + if not follower.cancelled() and follower.exception() is not None: + logger.warning( + "Live trajectory of turn %s stopped (continuing): %r", turn_id, follower.exception() + ) diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index 6b34136b..15642b08 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -7,6 +7,7 @@ import logging import time import uuid +from contextlib import nullcontext from datetime import datetime, timezone from typing import Any, ClassVar, Optional import httpx @@ -21,9 +22,11 @@ trajectory_mode, ) from agent_env.a2a_agent.staging import draining, staged_changelogs, transfer_store +from agent_env.a2a_agent.trajectory_follower import live_trajectory_endpoint from agent_env.config.model import MODEL_PARAMS_RESERVED from agent_env.env.gateway.constants import EXT_CLOCK_URI, EXT_TRIGGERS_URI, TRIGGER_IN_FLIGHT_STATUSES from agent_env.store import DuplicateKeyError, get_config +from agent_env.store.object_store import ObjectStore from agent_env.task_step.context import PromptResponse, TaskStepContext from agent_env.entity_refs import EntityRef from agent_env.task_step.task_step import RetryConfig, TaskStep, TaskStepDependency @@ -34,7 +37,9 @@ from agent_env.task_step.snapshot_utils.agent_state_capture import ( store_trajectory, trajectory_object_url, + trajectory_prefix_key, ) +from agent_env.task_step.snapshot_utils.live_trajectory import following_live_trajectory from agent_env.task_step.snapshot_utils.snapshot_series import SnapshotConfig, SnapshotSeries logger = logging.getLogger(__name__) @@ -147,6 +152,7 @@ def __init__( retry_config: Optional[RetryConfig | dict] = None, *, snapshot_config: Optional[SnapshotConfig] = None, + live_trajectory: bool = False, ): # `prompt` is the legacy text-only field; `parts` is the multimodal # form that's the authoritative A2A wire shape. Exactly one must be set: @@ -189,6 +195,7 @@ def __init__( self.user_output_format = user_output_format self.user_model = user_model self.snapshot_config = snapshot_config + self.live_trajectory = live_trajectory @staticmethod def _validate_parts(parts: list[dict]) -> None: @@ -262,6 +269,8 @@ def to_dict(self) -> dict: base["user_model"] = self.user_model if self.snapshot_config: base["snapshot_config"] = self.snapshot_config.to_dict() + if self.live_trajectory: + base["live_trajectory"] = True return base @classmethod @@ -297,6 +306,7 @@ def from_dict(cls, data: dict) -> PromptAgentTaskStep: if data.get("snapshot_config") else None ), + live_trajectory=data.get("live_trajectory", False), ) def _trajectory_prefix(self) -> str: @@ -477,6 +487,7 @@ async def _execute_conversation( per_turn_trajectory_s3_uris: list[str | None] = [] source_agent_per_turn_prompt_parts: list[list[dict] | None] = [] traj_ext_cached = A2AAgent.find_extension(card, A2AAgent.EXT_TRAJECTORY) + live = self._live_trajectory_target(context, card, target_url, trajectory_output_prefix) final_state: str = TaskState.completed.value trajectory_s3_uri: Optional[str] = None # What the target has been sent: of the files its replies name, the only ones a user-sim is made able @@ -509,11 +520,20 @@ def record_turn() -> None: else list(current_user_parts) ) - sent_task_id, result = await send_and_wait( - target_url, current_user_parts, agent=agent, message_id=target_a2a_task_id, - context_id=solver_context_id, timeout_seconds=self.timeout_seconds, - poll_interval_seconds=self.poll_interval_seconds, before_send=record_turn, + following = ( + following_live_trajectory( + *live, target_a2a_task_id, poll_interval_seconds=self.poll_interval_seconds + ) + if live is not None + else nullcontext(None) ) + async with following as on_sent: + sent_task_id, result = await send_and_wait( + target_url, current_user_parts, agent=agent, message_id=target_a2a_task_id, + context_id=solver_context_id, timeout_seconds=self.timeout_seconds, + poll_interval_seconds=self.poll_interval_seconds, before_send=record_turn, + on_sent=on_sent, + ) target_state = result["status"]["state"] status_msg = (result.get("status") or {}).get("message") or {} final_terminal = protocol.TerminalResponse.from_message(status_msg) @@ -719,6 +739,27 @@ def record_turn() -> None: return context + def _live_trajectory_target( + self, context: TaskStepContext, card: dict, target_url: str, trajectory_output_prefix: str + ) -> Optional[tuple[str, ObjectStore, str]]: + """Where this run follows each turn's trajectory live and stores it: the agent's live read, + the store and the prefix's key. None when the run did not ask for it, the agent does not + offer it, or the prefix is not in the configured store.""" + enabled = self.step_param_overrides(context).get("live_trajectory", self.live_trajectory) + if not bool(enabled): + return None + endpoint = live_trajectory_endpoint(card, target_url) + if endpoint is None: + logger.info(f"{self.id}: agent '{self.agent_name}' offers no live trajectory; not following it") + return None + store = get_config().get_object_store() + try: + prefix_key = trajectory_prefix_key(store, trajectory_output_prefix) + except ValueError as e: + logger.warning(f"{self.id}: not following the live trajectory: {e}") + return None + return endpoint, store, prefix_key + async def _fetch_trajectory( self, a2a_url: str, traj_ext: dict, a2a_server_task_id: str, trajectory_output_prefix: str, target_a2a_task_id: str, *, sandbox_type: str | None, card: dict | None = None, diff --git a/tst/unit/a2a_agent/test_trajectory_follower.py b/tst/unit/a2a_agent/test_trajectory_follower.py new file mode 100644 index 00000000..066d176c --- /dev/null +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -0,0 +1,206 @@ +import asyncio +import json +import time + +import httpx +import pytest +from agentenv_protocol.a2a_agent import TrajectoryState + +from agent_env.a2a_agent.trajectory_follower import ( + TrajectoryBatch, + follow_trajectory, + live_trajectory_endpoint, +) +from agent_env.store.object_store import LocalFilesystemObjectStore +from agent_env.task_step.snapshot_utils import live_trajectory +from agent_env.task_step.snapshot_utils.live_trajectory import ( + LiveTrajectoryChunks, + following_live_trajectory, +) + +_ENDPOINT = "http://agent.test/ext/trajectory" + + +def _page(after: int, events: list, state: str, *, has_more: bool = False) -> dict: + return { + "context_id": "ctx", + "task_id": "task-1", + "state": state, + "format": "test-events/1", + "events": events, + "next": after + len(events), + "has_more": has_more, + } + + +class _Agent: + """Answers each cursor read from a script of responses, recording the ``after`` it was sent.""" + + def __init__(self, *answers: httpx.Response) -> None: + self.answers = list(answers) + self.afters: list[int] = [] + + def client(self) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(self._answer)) + + def _answer(self, request: httpx.Request) -> httpx.Response: + self.afters.append(json.loads(request.content)["after"]) + return self.answers.pop(0) + + +async def _follow(agent: _Agent) -> list[TrajectoryBatch]: + batches: list[TrajectoryBatch] = [] + + async def sink(batch: TrajectoryBatch) -> None: + batches.append(batch) + + async with agent.client() as client: + await follow_trajectory( + _ENDPOINT, "task-1", sink, poll_interval_seconds=0, stop=asyncio.Event(), client=client + ) + return batches + + +def test_the_live_endpoint_comes_from_the_card(): + def card(request: dict) -> dict: + return { + "capabilities": { + "extensions": [ + { + "uri": "urn:agentenv:trajectory/v1", + "params": {"endpoint": "/ext/trajectory", "methods": {"get": {"method": "POST", "request": request}}}, + } + ] + } + } + + live = card({"oneOf": [{"required": ["task_id"]}, {"required": ["task_id", "after"], "optional": ["limit"]}]}) + assert live_trajectory_endpoint(live, "http://agent.test") == _ENDPOINT + assert live_trajectory_endpoint(card({"required": ["task_id"]}), "http://agent.test") is None + assert live_trajectory_endpoint({}, "http://agent.test") is None + + +@pytest.mark.asyncio +async def test_pages_are_followed_until_the_task_ends_and_every_event_is_read(): + agent = _Agent( + httpx.Response(200, json=_page(0, [], "pending")), + httpx.Response(200, json=_page(0, [{"i": 0}], "running")), + httpx.Response(200, json=_page(1, [{"i": 1}], "completed", has_more=True)), + httpx.Response(200, json=_page(2, [{"i": 2}], "completed")), + ) + + batches = await _follow(agent) + + assert agent.afters == [0, 0, 1, 2] + assert [(b.after, b.next, b.events, b.last) for b in batches] == [ + (0, 1, [{"i": 0}], False), + (1, 2, [{"i": 1}], False), + (2, 3, [{"i": 2}], True), + ] + assert batches[-1].state is TrajectoryState.COMPLETED + + +@pytest.mark.asyncio +async def test_an_ended_task_with_nothing_new_still_gets_a_last_batch(): + agent = _Agent(httpx.Response(200, json=_page(0, [], "failed"))) + + batches = await _follow(agent) + + assert [(b.events, b.state, b.last) for b in batches] == [([], TrajectoryState.FAILED, True)] + + +@pytest.mark.asyncio +async def test_a_server_error_is_retried(monkeypatch): + monkeypatch.setattr("agent_env.a2a_agent.trajectory_follower._MAX_BACKOFF_SECONDS", 0) + agent = _Agent( + httpx.Response(503), + httpx.Response(200, json=_page(0, [{"i": 0}], "completed")), + ) + + batches = await _follow(agent) + + assert agent.afters == [0, 0] + assert batches[-1].events == [{"i": 0}] + + +@pytest.mark.asyncio +async def test_a_task_the_agent_no_longer_knows_ends_following(): + agent = _Agent(httpx.Response(404, json={"detail": "Unknown task"})) + + assert await _follow(agent) == [] + + +@pytest.mark.asyncio +async def test_a_client_error_is_raised(): + agent = _Agent(httpx.Response(400, json={"detail": "bad"})) + + with pytest.raises(httpx.HTTPStatusError): + await _follow(agent) + + +@pytest.mark.asyncio +async def test_chunks_meta_and_end_are_stored_once(tmp_path): + store = LocalFilesystemObjectStore(str(tmp_path)) + chunks = LiveTrajectoryChunks(store, "prefix/", "turn") + batch = TrajectoryBatch( + after=0, next=2, events=[{"i": 0}, {"i": 1}], format="test-events/1", + state=TrajectoryState.COMPLETED, last=True, + ) + + await chunks(batch) + await LiveTrajectoryChunks(store, "prefix/", "turn")(batch) + + assert sorted(store.list("prefix/")) == [ + "prefix/turn/live/00000000-00000002.jsonl", + "prefix/turn/live/end.json", + "prefix/turn/live/meta.json", + ] + assert store.get(store.object_url("prefix/turn/live/00000000-00000002.jsonl")) == b'{"i":0}\n{"i":1}\n' + assert json.loads(store.get(store.object_url("prefix/turn/live/end.json"))) == {"state": "completed", "next": 2} + + +@pytest.mark.asyncio +async def test_a_failing_write_is_retried_then_raised(monkeypatch, tmp_path): + store = LocalFilesystemObjectStore(str(tmp_path)) + attempts: list[str] = [] + + def failing_put(key, body, *, content_type): + attempts.append(key) + raise OSError("disk full") + + async def no_sleep(_seconds): + return None + + monkeypatch.setattr(store, "put", failing_put) + monkeypatch.setattr(live_trajectory.asyncio, "sleep", no_sleep) + batch = TrajectoryBatch(after=0, next=1, events=[{"i": 0}], format=None, state=TrajectoryState.RUNNING, last=False) + + with pytest.raises(OSError): + await LiveTrajectoryChunks(store, "prefix/", "turn")(batch) + assert len(attempts) == live_trajectory._WRITE_ATTEMPTS + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [RuntimeError("turn failed"), asyncio.CancelledError()]) +async def test_a_failed_turn_stops_its_follower_without_draining(monkeypatch, tmp_path, error): + """An agent that never ends: only a drain would keep the follower reading.""" + monkeypatch.setattr( + httpx, + "AsyncClient", + lambda *a, real=httpx.AsyncClient, **kw: real( + *a, transport=httpx.MockTransport(lambda r: httpx.Response(200, json=_page(0, [], "running"))), **kw + ), + ) + store = LocalFilesystemObjectStore(str(tmp_path)) + started = time.monotonic() + + with pytest.raises(type(error)): + async with following_live_trajectory( + _ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0 + ) as on_sent: + on_sent("task-1") + await asyncio.sleep(0.05) + raise error + + assert time.monotonic() - started < 2 + assert store.list("prefix/") == [] diff --git a/tst/unit/task_step/test_prompt_agent_live_trajectory.py b/tst/unit/task_step/test_prompt_agent_live_trajectory.py new file mode 100644 index 00000000..091e7fe2 --- /dev/null +++ b/tst/unit/task_step/test_prompt_agent_live_trajectory.py @@ -0,0 +1,174 @@ +"""prompt_agent following a turn's trajectory live and storing it in chunks, against a live SDK +agent served in-process.""" + +import asyncio +import json +import logging + +import httpx +import pytest +from agentenv_protocol.a2a_agent import ( + TRAJECTORY_V1, + AgentEnvAgent, + AgentIdentity, + TaskRequest, + TaskResult, + a2a_agent, + enable, +) + +from agent_env.config import configure +from agent_env.store.object_store import LocalFilesystemObjectStore +from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.task_steps import prompt_agent as pa +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep + +_PREFIX = "prompt_agent_trajectories/prompt_id=solve/" +_EVENTS = [{"type": "system", "subtype": "init"}, {"type": "assistant", "step": 0}] +_MORE = [{"type": "assistant", "step": 1}, {"type": "result", "result": "done"}] + + +def _agent(store: LocalFilesystemObjectStore, *, live: bool, wait_for_chunk: bool) -> AgentEnvAgent: + """Appends two events, optionally waits until a chunk of them is stored, then two more.""" + + @a2a_agent( + identity=AgentIdentity(name="solver", description="test", version="1"), + extensions=(enable(TRAJECTORY_V1, live=live),), + ) + class Agent(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.set_format("test-events/1") + for event in _EVENTS: + request.trajectory.append(event) + if wait_for_chunk: + async with asyncio.timeout(5): + while not _chunk_keys(store): + await asyncio.sleep(0.01) + for event in _MORE: + request.trajectory.append(event) + return TaskResult.text("done") + + return Agent() + + +def _chunk_keys(store: LocalFilesystemObjectStore) -> list[str]: + return sorted(key for key in store.list(_PREFIX) if key.endswith(".jsonl")) + + +class _Run: + def __init__(self, monkeypatch: pytest.MonkeyPatch, tmp_path, *, live: bool, wait_for_chunk: bool) -> None: + self.store = LocalFilesystemObjectStore(str(tmp_path)) + configure(object_store=self.store) + for fn in ("create_conversation", "add_a2a_task", "complete_a2a_task", "mark_closed", "get_conversation"): + monkeypatch.setattr(pa.conversation_store, fn, lambda *a, **kw: None) + app = _agent(self.store, live=live, wait_for_chunk=wait_for_chunk).create_app() + self.card = app.state.agentenv_a2a.card.model_dump(mode="json", exclude_none=True) + self.cursor_reads: list[dict] = [] + transport = httpx.ASGITransport(app=app) + real_client = httpx.AsyncClient + + async def record(request: httpx.Request) -> None: + if request.url.path == "/ext/trajectory": + body = json.loads(await request.aread()) + if "after" in body: + self.cursor_reads.append(body) + + monkeypatch.setattr( + httpx, + "AsyncClient", + lambda *a, **kw: real_client(*a, transport=transport, event_hooks={"request": [record]}, **kw), + ) + + async def execute(self, step: PromptAgentTaskStep, context: TaskStepContext | None = None) -> TaskStepContext: + context = context or TaskStepContext(instance_id="ti-1") + context.deployed_agents.append( + DeployedAgent(agent_name="solver", api_url="http://agent.test", a2a_url="http://agent.test", a2a_card=self.card) + ) + return await step.execute(context) + + def step(self, **kwargs) -> PromptAgentTaskStep: + return PromptAgentTaskStep( + id="solve", version=None, prompt="go", agent_name="solver", poll_interval_seconds=0, + trajectory_output_prefix=self.store.object_url(_PREFIX), **kwargs, + ) + + +@pytest.mark.asyncio +async def test_a_live_turn_is_stored_in_chunks_while_it_runs(monkeypatch, tmp_path): + run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=True) + + result = await run.execute(run.step(live_trajectory=True)) + + final_url = result.prompt_responses[-1].agent_trajectory_s3_uri + turn_id = run.store.get_object_key(final_url).removeprefix(f"{_PREFIX}trajectory-").removesuffix(".json") + live = f"{_PREFIX}{turn_id}/live/" + chunks = _chunk_keys(run.store) + assert chunks and all(key.startswith(live) for key in chunks) + assert chunks[0] == f"{live}00000000-{len(_EVENTS):08d}.jsonl" + joined = [ + json.loads(line) + for key in chunks + for line in run.store.get(run.store.object_url(key)).decode().splitlines() + ] + assert joined == _EVENTS + _MORE == json.loads(run.store.get(final_url)) + assert json.loads(run.store.get(run.store.object_url(f"{live}meta.json"))) == {"format": "test-events/1"} + assert json.loads(run.store.get(run.store.object_url(f"{live}end.json"))) == { + "state": "completed", + "next": len(_EVENTS + _MORE), + } + + +@pytest.mark.asyncio +async def test_without_live_trajectory_nothing_is_followed(monkeypatch, tmp_path): + run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=False) + + result = await run.execute(run.step()) + + assert run.cursor_reads == [] + assert _chunk_keys(run.store) == [] + assert json.loads(run.store.get(result.prompt_responses[-1].agent_trajectory_s3_uri)) == _EVENTS + _MORE + + +@pytest.mark.asyncio +async def test_a_step_params_override_turns_it_on(monkeypatch, tmp_path): + run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=True) + context = TaskStepContext(instance_id="ti-1") + context.metadata["user_overrides"] = {"step_params": {"solve": {"live_trajectory": True}}} + + await run.execute(run.step(), context) + + assert _chunk_keys(run.store) + + +@pytest.mark.asyncio +async def test_an_agent_without_the_live_read_is_not_followed(monkeypatch, tmp_path): + run = _Run(monkeypatch, tmp_path, live=False, wait_for_chunk=False) + + result = await run.execute(run.step(live_trajectory=True)) + + assert run.cursor_reads == [] + assert _chunk_keys(run.store) == [] + assert json.loads(run.store.get(result.prompt_responses[-1].agent_trajectory_s3_uri)) == _EVENTS + _MORE + + +def test_a_prefix_outside_the_store_is_not_followed(monkeypatch, tmp_path, caplog): + run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=False) + step = PromptAgentTaskStep(id="solve", version=None, prompt="go", live_trajectory=True) + + with caplog.at_level(logging.WARNING): + target = step._live_trajectory_target( + TaskStepContext(), run.card, "http://agent.test", "s3://another-bucket/prefix/" + ) + + assert target is None + assert "not following the live trajectory" in caplog.text + + +def test_live_trajectory_is_serialized_only_when_on(): + off = PromptAgentTaskStep(id="solve", version=None, prompt="go") + on = PromptAgentTaskStep(id="solve", version=None, prompt="go", live_trajectory=True) + + assert "live_trajectory" not in off.to_dict() + assert on.to_dict()["live_trajectory"] is True + assert PromptAgentTaskStep.from_dict(on.to_dict()).live_trajectory is True + assert PromptAgentTaskStep.from_dict(off.to_dict()).live_trajectory is False From 367ed3bfb2427b3c3c12f75fec1cb9ef67a7e453 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 7 Oct 2026 21:52:31 -0700 Subject: [PATCH 05/12] fix(prompt_agent): mark a stopped turn's live trajectory ended When the step stops following before the agent's task ends (a cancelled step, or a turn whose send/wait raised), the follower is stopped at once and never reads a terminal state, so `end.json` was never written and readers could not tell the turn was over: AMV reported a cancelled run's turn as live. `following_live_trajectory` now writes `end.json` itself in that case, `{"state": "canceled" | "failed", "next": }`, bounded to 5 s so the cancel is not held up. `end.json` is write-once, so a state the follower already read (e.g. `completed`) stays. FD-3719 Co-Authored-By: Claude Opus 5.5 --- .../snapshot_utils/live_trajectory.py | 50 +++++++++++-- .../a2a_agent/test_trajectory_follower.py | 71 ++++++++++++++++--- 2 files changed, 108 insertions(+), 13 deletions(-) diff --git a/src/agent_env/task_step/snapshot_utils/live_trajectory.py b/src/agent_env/task_step/snapshot_utils/live_trajectory.py index 777e4088..046db625 100644 --- a/src/agent_env/task_step/snapshot_utils/live_trajectory.py +++ b/src/agent_env/task_step/snapshot_utils/live_trajectory.py @@ -1,5 +1,7 @@ """A turn's live trajectory, stored as it is read: ``/live/-.jsonl`` per -page, ``meta.json`` naming the events' format, and ``end.json`` once the turn has ended. +page, ``meta.json`` naming the events' format, and ``end.json`` once the turn has ended: with the +agent's final state when the follower read it, else ``canceled`` or ``failed`` when the step stopped +following first. Written from the worker with its own credentials, like the final trajectory. Every object is written once: one already there holds the same bytes, so a repeated write is a no-op.""" @@ -11,6 +13,8 @@ from collections.abc import AsyncIterator, Callable from contextlib import asynccontextmanager +from agentenv_protocol.a2a_agent import TrajectoryState + from agent_env.a2a_agent.trajectory_follower import TrajectoryBatch, follow_trajectory from agent_env.store.base import ObjectAlreadyExistsError from agent_env.store.object_store import ObjectStore @@ -20,6 +24,8 @@ _WRITE_ATTEMPTS = 3 # The task has ended when the turn's wait returns, so draining its last events is normally one read. _DRAIN_SECONDS = 60 +# A stopped step marks its turn ended in passing: the write must not hold up the cancel. +_END_WRITE_SECONDS = 5 class LiveTrajectoryChunks: @@ -29,6 +35,8 @@ def __init__(self, store: ObjectStore, prefix_key: str, turn_id: str) -> None: self._store = store self._directory = f"{prefix_key}{turn_id}/live/" self._format_written = False + self._stored = 0 + self._ended = False async def __call__(self, batch: TrajectoryBatch) -> None: if batch.events: @@ -36,12 +44,20 @@ async def __call__(self, batch: TrajectoryBatch) -> None: json.dumps(event, separators=(",", ":")).encode() + b"\n" for event in batch.events ) await self._put(f"{batch.after:08d}-{batch.next:08d}.jsonl", body, "application/x-ndjson") + self._stored = batch.next if batch.format and not self._format_written: await self._put("meta.json", json.dumps({"format": batch.format}).encode(), "application/json") self._format_written = True if batch.last: - end = {"state": batch.state.value, "next": batch.next} - await self._put("end.json", json.dumps(end).encode(), "application/json") + await self.end(batch.state) + + async def end(self, state: TrajectoryState) -> None: + """Mark the turn ended in ``state`` after the events stored so far; the first mark stays.""" + if self._ended: + return + end = {"state": state.value, "next": self._stored} + await self._put("end.json", json.dumps(end).encode(), "application/json") + self._ended = True async def _put(self, name: str, body: bytes, content_type: str) -> None: key = self._directory + name @@ -70,13 +86,17 @@ async def following_live_trajectory( """Yield the callback that starts following the agent's task once the turn's message is sent. When the body returns, the follower drains what remains, for at most ``_DRAIN_SECONDS``; when it - raises, the follower is stopped at once, since the agent may be gone. A follower's own failure - is logged and never fails the turn.""" + raises, the follower is stopped at once, since the agent may be gone, and the turn is marked + ended (``canceled`` for a cancelled step, else ``failed``) so readers do not wait on it. A + follower's own failure is logged and never fails the turn.""" stop = asyncio.Event() followers: list[asyncio.Task] = [] + sinks: list[LiveTrajectoryChunks] = [] + stopped_as: TrajectoryState | None = None def follow(task_id: str) -> None: sink = LiveTrajectoryChunks(store, prefix_key, turn_id) + sinks.append(sink) followers.append( asyncio.create_task( follow_trajectory( @@ -96,6 +116,12 @@ def follow(task_id: str) -> None: turn_id, _DRAIN_SECONDS, ) + except asyncio.CancelledError: + stopped_as = TrajectoryState.CANCELED + raise + except BaseException: + stopped_as = TrajectoryState.FAILED + raise finally: stop.set() for follower in followers: @@ -107,3 +133,17 @@ def follow(task_id: str) -> None: logger.warning( "Live trajectory of turn %s stopped (continuing): %r", turn_id, follower.exception() ) + if stopped_as is not None and sinks: + await _mark_ended(sinks[-1], stopped_as, turn_id) + + +async def _mark_ended(sink: LiveTrajectoryChunks, state: TrajectoryState, turn_id: str) -> None: + marking = asyncio.create_task(sink.end(state)) + done, _ = await asyncio.wait({marking}, timeout=_END_WRITE_SECONDS) + if not done: + marking.cancel() + logger.warning("Live trajectory of turn %s was not marked %s in time (continuing)", turn_id, state.value) + elif marking.exception() is not None: + logger.warning( + "Live trajectory of turn %s was not marked %s (continuing): %r", turn_id, state.value, marking.exception() + ) diff --git a/tst/unit/a2a_agent/test_trajectory_follower.py b/tst/unit/a2a_agent/test_trajectory_follower.py index 066d176c..ab67f195 100644 --- a/tst/unit/a2a_agent/test_trajectory_follower.py +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -180,17 +180,35 @@ async def no_sleep(_seconds): assert len(attempts) == live_trajectory._WRITE_ATTEMPTS -@pytest.mark.asyncio -@pytest.mark.parametrize("error", [RuntimeError("turn failed"), asyncio.CancelledError()]) -async def test_a_failed_turn_stops_its_follower_without_draining(monkeypatch, tmp_path, error): - """An agent that never ends: only a drain would keep the follower reading.""" +def _agent_answering(monkeypatch, *pages: dict) -> None: + """Every httpx client reaches an agent that answers each cursor read with the next page, then + keeps answering with the last one.""" + answers = list(pages) + + def answer(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=answers.pop(0) if len(answers) > 1 else answers[0]) + monkeypatch.setattr( httpx, "AsyncClient", - lambda *a, real=httpx.AsyncClient, **kw: real( - *a, transport=httpx.MockTransport(lambda r: httpx.Response(200, json=_page(0, [], "running"))), **kw - ), + lambda *a, real=httpx.AsyncClient, **kw: real(*a, transport=httpx.MockTransport(answer), **kw), ) + + +def _end(store: LocalFilesystemObjectStore) -> dict: + return json.loads(store.get(store.object_url("prefix/turn/live/end.json"))) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error", "state"), + [(RuntimeError("turn failed"), "failed"), (asyncio.CancelledError(), "canceled")], +) +async def test_a_failed_turn_stops_its_follower_without_draining_and_marks_it_ended( + monkeypatch, tmp_path, error, state +): + """An agent that never ends: only a drain would keep the follower reading.""" + _agent_answering(monkeypatch, _page(0, [], "running")) store = LocalFilesystemObjectStore(str(tmp_path)) started = time.monotonic() @@ -203,4 +221,41 @@ async def test_a_failed_turn_stops_its_follower_without_draining(monkeypatch, tm raise error assert time.monotonic() - started < 2 - assert store.list("prefix/") == [] + assert store.list("prefix/") == ["prefix/turn/live/end.json"] + assert _end(store) == {"state": state, "next": 0} + + +@pytest.mark.asyncio +async def test_a_cancelled_turn_is_marked_ended_after_the_events_it_stored(monkeypatch, tmp_path): + _agent_answering(monkeypatch, _page(0, [{"i": 0}, {"i": 1}], "running"), _page(2, [], "running")) + store = LocalFilesystemObjectStore(str(tmp_path)) + + with pytest.raises(asyncio.CancelledError): + async with following_live_trajectory( + _ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0 + ) as on_sent: + on_sent("task-1") + async with asyncio.timeout(5): + while "prefix/turn/live/00000000-00000002.jsonl" not in store.list("prefix/"): + await asyncio.sleep(0.01) + raise asyncio.CancelledError() + + assert _end(store) == {"state": "canceled", "next": 2} + + +@pytest.mark.asyncio +async def test_a_turn_the_follower_saw_end_keeps_its_state(monkeypatch, tmp_path): + _agent_answering(monkeypatch, _page(0, [{"i": 0}], "completed")) + store = LocalFilesystemObjectStore(str(tmp_path)) + + with pytest.raises(RuntimeError): + async with following_live_trajectory( + _ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0 + ) as on_sent: + on_sent("task-1") + async with asyncio.timeout(5): + while "prefix/turn/live/end.json" not in store.list("prefix/"): + await asyncio.sleep(0.01) + raise RuntimeError("the step failed after the turn ended") + + assert _end(store) == {"state": "completed", "next": 1} From 0b3c84c08298f55dbe1c2e3d4b36d627381e2268 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Wed, 7 Oct 2026 21:50:34 -0700 Subject: [PATCH 06/12] test(protocol): the log's final trajectory is byte-identical to a native one The same events, returned as a `native_trajectory` by one agent and appended to the log by another, give the same `{task_id}` answer and the same uploaded bytes (non-ASCII text, U+2028, nested objects, floats and nulls included), so porting an agent to the log does not change its final trajectory. FD-3719 Co-Authored-By: Claude Opus 5.5 --- .../tests/test_a2a_trajectory_log.py | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py index 8cd1d4ad..84327bbb 100644 --- a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py +++ b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py @@ -5,6 +5,7 @@ import logging import threading import time +from datetime import datetime, timedelta, timezone from typing import Any, AsyncIterator import pytest @@ -24,12 +25,14 @@ TaskTrajectoryRequest, TrajectoryLog, TrajectoryState, + Uploaded, a2a_agent, build_registry, card_request_accepts, enable, extension, ) +from agentenv_protocol.a2a_agent import framework from agentenv_protocol.a2a_agent._trajectory_log import TaskTrajectories from agentenv_protocol.a2a_agent.framework import _SdkServices from agentenv_protocol.a2a_agent.tasks import v1 as tasks_v1 @@ -552,3 +555,68 @@ def test_a_context_read_keeps_the_format_while_a_new_turn_is_pending() -> None: TrajectoryState.PENDING, "test-events/1", ) + + +# Events shaped like a CLI's: nested objects, non-ASCII and line-separator text, floats, nulls. +_NATIVE_EVENTS = [ + {"type": "system", "subtype": "init", "tools": ["Read", "mcp__db__query"], "cwd": "/app"}, + {"type": "assistant", "message": {"content": [{"type": "text", "text": "Prüfe die Daten — ✓ a\u2028b"}]}}, + {"type": "user", "message": {"content": [{"type": "tool_result", "content": "ok", "is_error": False}]}}, + {"type": "result", "subtype": "success", "result": "done", "total_cost_usd": 0.0123, "usage": None}, +] + + +def test_the_log_is_byte_identical_to_the_native_trajectory_it_replaces(monkeypatch): + """An agent that appends its events and one that returns them as a native trajectory answer the + task read with the same trajectory and upload the same bytes, so porting an agent to the log does + not change its final trajectory.""" + + @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1),)) + class Native(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + return ( + TaskResult.builder() + .succeeded() + .add_text("done") + .native_trajectory(format="test-events/1", payload=_NATIVE_EVENTS) + .build() + ) + + @a2a_agent(identity=_IDENTITY, extensions=(enable(TRAJECTORY_V1, live=True),)) + class Live(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + request.trajectory.set_format("test-events/1") + for event in _NATIVE_EVENTS: + request.trajectory.append(event) + return TaskResult.text("done") + + uploaded: list[bytes] = [] + + async def capture_upload(_target, body: bytes): + uploaded.append(body) + return Uploaded(size_bytes=len(body)) + + monkeypatch.setattr(framework, "upload", capture_upload) + grant = { + "media_type": "application/json", + "max_bytes": 1_000_000, + "write": { + "kind": "http-put", + "url": "https://objects.example.test/write", + "expires_at": (datetime.now(timezone.utc) + timedelta(hours=1)) + .isoformat(timespec="seconds") + .replace("+00:00", "Z"), + }, + } + inline: list[Any] = [] + for agent in (Native(), Live()): + with TestClient(agent.create_app()) as client: + task_id = _send(client, "go", "message-1") + assert _within(5, lambda: _task_state(client, task_id) == "completed") + inline.append(_trajectory(client, {"task_id": task_id}).json()) + response = _trajectory(client, {"task_id": task_id, "objects": {"trajectory": grant}}) + assert response.status_code == 200, response.text + + assert inline[0] == inline[1] == {"trajectory": _NATIVE_EVENTS} + assert uploaded[0] == uploaded[1] + assert json.loads(uploaded[1]) == _NATIVE_EVENTS From 7118de8afd6b9cb757f2ac9d164c0d1f2d4b5222 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 8 Oct 2026 15:25:58 -0700 Subject: [PATCH 07/12] fix(prompt_agent): mark a turn's live trajectory ended whenever following stops A turn whose step returned normally got no end.json when the follower died on a client error, the agent no longer knew the task, or the drain timed out, so readers kept waiting on it. The step now reports the state its task ended in, and the turn is marked ended in it whenever the follower did not read the end. Co-Authored-By: Claude Opus 5.5 --- .../snapshot_utils/live_trajectory.py | 97 +++++++++++-------- .../task_step/task_steps/prompt_agent.py | 8 +- .../a2a_agent/test_trajectory_follower.py | 66 ++++++++++--- .../test_prompt_agent_live_trajectory.py | 21 ++++ 4 files changed, 140 insertions(+), 52 deletions(-) diff --git a/src/agent_env/task_step/snapshot_utils/live_trajectory.py b/src/agent_env/task_step/snapshot_utils/live_trajectory.py index 046db625..1d0e82dc 100644 --- a/src/agent_env/task_step/snapshot_utils/live_trajectory.py +++ b/src/agent_env/task_step/snapshot_utils/live_trajectory.py @@ -1,7 +1,8 @@ """A turn's live trajectory, stored as it is read: ``/live/-.jsonl`` per -page, ``meta.json`` naming the events' format, and ``end.json`` once the turn has ended: with the -agent's final state when the follower read it, else ``canceled`` or ``failed`` when the step stopped -following first. +page, ``meta.json`` naming the events' format, and ``end.json`` once following stops: with the +agent's final state when the follower read it, else the state the step saw the turn end in, or +``canceled`` or ``failed`` when the step stopped first. A follower that stopped early leaves +``end.json``'s ``next`` short of the turn's events; the final trajectory holds them all. Written from the worker with its own credentials, like the final trajectory. Every object is written once: one already there holds the same bytes, so a repeated write is a no-op.""" @@ -10,7 +11,7 @@ import asyncio import json import logging -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator from contextlib import asynccontextmanager from agentenv_protocol.a2a_agent import TrajectoryState @@ -26,6 +27,8 @@ _DRAIN_SECONDS = 60 # A stopped step marks its turn ended in passing: the write must not hold up the cancel. _END_WRITE_SECONDS = 5 +# The A2A task states a turn's trajectory ends in as themselves; any other (failed, rejected) ends it failed. +_ENDED_AS = {"completed": TrajectoryState.COMPLETED, "canceled": TrajectoryState.CANCELED} class LiveTrajectoryChunks: @@ -74,6 +77,39 @@ async def _put(self, name: str, body: bytes, content_type: str) -> None: await asyncio.sleep(attempt) +class LiveTurn: + """What a turn's body tells its follower: the task to ``follow`` once the message is sent, and + the state it ``ended`` in once the wait for it returns.""" + + def __init__( + self, endpoint: str, store: ObjectStore, prefix_key: str, turn_id: str, poll_interval_seconds: float + ) -> None: + self._endpoint = endpoint + self._store = store + self._prefix_key = prefix_key + self._turn_id = turn_id + self._poll_interval_seconds = poll_interval_seconds + self.stop = asyncio.Event() + self.followers: list[asyncio.Task] = [] + self.sinks: list[LiveTrajectoryChunks] = [] + self.ended_as: TrajectoryState | None = None + + def follow(self, task_id: str) -> None: + sink = LiveTrajectoryChunks(self._store, self._prefix_key, self._turn_id) + self.sinks.append(sink) + self.followers.append( + asyncio.create_task( + follow_trajectory( + self._endpoint, task_id, sink, poll_interval_seconds=self._poll_interval_seconds, stop=self.stop + ) + ) + ) + + def ended(self, task_state: str) -> None: + """Report the A2A state the turn's task ended in.""" + self.ended_as = _ENDED_AS.get(task_state, TrajectoryState.FAILED) + + @asynccontextmanager async def following_live_trajectory( endpoint: str, @@ -82,34 +118,22 @@ async def following_live_trajectory( turn_id: str, *, poll_interval_seconds: float, -) -> AsyncIterator[Callable[[str], None]]: - """Yield the callback that starts following the agent's task once the turn's message is sent. +) -> AsyncIterator[LiveTurn]: + """Yield the turn whose task is followed once its message is sent. When the body returns, the follower drains what remains, for at most ``_DRAIN_SECONDS``; when it - raises, the follower is stopped at once, since the agent may be gone, and the turn is marked - ended (``canceled`` for a cancelled step, else ``failed``) so readers do not wait on it. A - follower's own failure is logged and never fails the turn.""" - stop = asyncio.Event() - followers: list[asyncio.Task] = [] - sinks: list[LiveTrajectoryChunks] = [] - stopped_as: TrajectoryState | None = None - - def follow(task_id: str) -> None: - sink = LiveTrajectoryChunks(store, prefix_key, turn_id) - sinks.append(sink) - followers.append( - asyncio.create_task( - follow_trajectory( - endpoint, task_id, sink, poll_interval_seconds=poll_interval_seconds, stop=stop - ) - ) - ) + raises, the follower is stopped at once, since the agent may be gone. Either way, a turn whose + end the follower did not read is marked ended in the state the body reported, else ``canceled`` + for a cancelled step or ``failed``, so readers do not wait on it. A follower's own failure is + logged and never fails the turn.""" + turn = LiveTurn(endpoint, store, prefix_key, turn_id, poll_interval_seconds) + stopped_as = TrajectoryState.FAILED try: - yield follow - stop.set() - if followers: - _, draining = await asyncio.wait(followers, timeout=_DRAIN_SECONDS) + yield turn + turn.stop.set() + if turn.followers: + _, draining = await asyncio.wait(turn.followers, timeout=_DRAIN_SECONDS) if draining: logger.warning( "Live trajectory of turn %s did not finish draining in %ds (continuing)", @@ -119,22 +143,19 @@ def follow(task_id: str) -> None: except asyncio.CancelledError: stopped_as = TrajectoryState.CANCELED raise - except BaseException: - stopped_as = TrajectoryState.FAILED - raise finally: - stop.set() - for follower in followers: + turn.stop.set() + for follower in turn.followers: follower.cancel() - if followers: - await asyncio.wait(followers) - for follower in followers: + if turn.followers: + await asyncio.wait(turn.followers) + for follower in turn.followers: if not follower.cancelled() and follower.exception() is not None: logger.warning( "Live trajectory of turn %s stopped (continuing): %r", turn_id, follower.exception() ) - if stopped_as is not None and sinks: - await _mark_ended(sinks[-1], stopped_as, turn_id) + if turn.sinks: + await _mark_ended(turn.sinks[-1], turn.ended_as or stopped_as, turn_id) async def _mark_ended(sink: LiveTrajectoryChunks, state: TrajectoryState, turn_id: str) -> None: diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index 15642b08..882a8714 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -527,14 +527,16 @@ def record_turn() -> None: if live is not None else nullcontext(None) ) - async with following as on_sent: + async with following as live_turn: sent_task_id, result = await send_and_wait( target_url, current_user_parts, agent=agent, message_id=target_a2a_task_id, context_id=solver_context_id, timeout_seconds=self.timeout_seconds, poll_interval_seconds=self.poll_interval_seconds, before_send=record_turn, - on_sent=on_sent, + on_sent=live_turn.follow if live_turn else None, ) - target_state = result["status"]["state"] + target_state = result["status"]["state"] + if live_turn: + live_turn.ended(target_state) status_msg = (result.get("status") or {}).get("message") or {} final_terminal = protocol.TerminalResponse.from_message(status_msg) agent_response_parts = status_msg.get("parts") or [{"kind": "text", "text": final_terminal.response_text}] diff --git a/tst/unit/a2a_agent/test_trajectory_follower.py b/tst/unit/a2a_agent/test_trajectory_follower.py index ab67f195..02627a6e 100644 --- a/tst/unit/a2a_agent/test_trajectory_follower.py +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -180,13 +180,14 @@ async def no_sleep(_seconds): assert len(attempts) == live_trajectory._WRITE_ATTEMPTS -def _agent_answering(monkeypatch, *pages: dict) -> None: - """Every httpx client reaches an agent that answers each cursor read with the next page, then - keeps answering with the last one.""" - answers = list(pages) +def _agent_answering(monkeypatch, *answers: dict | httpx.Response) -> None: + """Every httpx client reaches an agent that answers each cursor read with the next answer, a + page or a response, then keeps answering with the last one.""" + remaining = list(answers) def answer(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, json=answers.pop(0) if len(answers) > 1 else answers[0]) + given = remaining.pop(0) if len(remaining) > 1 else remaining[0] + return given if isinstance(given, httpx.Response) else httpx.Response(200, json=given) monkeypatch.setattr( httpx, @@ -215,8 +216,8 @@ async def test_a_failed_turn_stops_its_follower_without_draining_and_marks_it_en with pytest.raises(type(error)): async with following_live_trajectory( _ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0 - ) as on_sent: - on_sent("task-1") + ) as turn: + turn.follow("task-1") await asyncio.sleep(0.05) raise error @@ -233,8 +234,8 @@ async def test_a_cancelled_turn_is_marked_ended_after_the_events_it_stored(monke with pytest.raises(asyncio.CancelledError): async with following_live_trajectory( _ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0 - ) as on_sent: - on_sent("task-1") + ) as turn: + turn.follow("task-1") async with asyncio.timeout(5): while "prefix/turn/live/00000000-00000002.jsonl" not in store.list("prefix/"): await asyncio.sleep(0.01) @@ -251,11 +252,54 @@ async def test_a_turn_the_follower_saw_end_keeps_its_state(monkeypatch, tmp_path with pytest.raises(RuntimeError): async with following_live_trajectory( _ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0 - ) as on_sent: - on_sent("task-1") + ) as turn: + turn.follow("task-1") async with asyncio.timeout(5): while "prefix/turn/live/end.json" not in store.list("prefix/"): await asyncio.sleep(0.01) raise RuntimeError("the step failed after the turn ended") assert _end(store) == {"state": "completed", "next": 1} + + +async def _until_stored(store: LocalFilesystemObjectStore, key: str) -> None: + async with asyncio.timeout(5): + while key not in store.list("prefix/"): + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stops_with", + [ + pytest.param(httpx.Response(400, json={"detail": "bad"}), id="client-error"), + pytest.param(httpx.Response(404, json={"detail": "Unknown task"}), id="task-evicted"), + pytest.param(_page(2, [], "running"), id="drain-timeout"), + ], +) +async def test_a_turn_the_follower_stopped_before_its_end_is_marked_ended_in_the_reported_state( + monkeypatch, tmp_path, stops_with +): + monkeypatch.setattr(live_trajectory, "_DRAIN_SECONDS", 0.1) + _agent_answering(monkeypatch, _page(0, [{"i": 0}, {"i": 1}], "running"), stops_with) + store = LocalFilesystemObjectStore(str(tmp_path)) + + async with following_live_trajectory(_ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0) as turn: + turn.follow("task-1") + await _until_stored(store, "prefix/turn/live/00000000-00000002.jsonl") + turn.ended("completed") + + assert _end(store) == {"state": "completed", "next": 2} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("task_state", "state"), [("canceled", "canceled"), ("failed", "failed"), ("rejected", "failed")]) +async def test_a_reported_task_state_ends_the_turn_as_its_trajectory_state(monkeypatch, tmp_path, task_state, state): + _agent_answering(monkeypatch, httpx.Response(404, json={"detail": "Unknown task"})) + store = LocalFilesystemObjectStore(str(tmp_path)) + + async with following_live_trajectory(_ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0) as turn: + turn.follow("task-1") + turn.ended(task_state) + + assert _end(store) == {"state": state, "next": 0} diff --git a/tst/unit/task_step/test_prompt_agent_live_trajectory.py b/tst/unit/task_step/test_prompt_agent_live_trajectory.py index 091e7fe2..84111bad 100644 --- a/tst/unit/task_step/test_prompt_agent_live_trajectory.py +++ b/tst/unit/task_step/test_prompt_agent_live_trajectory.py @@ -20,6 +20,7 @@ from agent_env.config import configure from agent_env.store.object_store import LocalFilesystemObjectStore from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.snapshot_utils import live_trajectory from agent_env.task_step.task_steps import prompt_agent as pa from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep @@ -118,6 +119,26 @@ async def test_a_live_turn_is_stored_in_chunks_while_it_runs(monkeypatch, tmp_pa } +@pytest.mark.asyncio +async def test_a_turn_whose_follower_stopped_early_is_marked_ended_in_the_turns_state(monkeypatch, tmp_path): + async def stops_at_once(*_args, **_kwargs) -> None: + return None + + monkeypatch.setattr(live_trajectory, "follow_trajectory", stops_at_once) + run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=False) + + result = await run.execute(run.step(live_trajectory=True)) + + final_key = run.store.get_object_key(result.prompt_responses[-1].agent_trajectory_s3_uri) + turn_id = final_key.removeprefix(f"{_PREFIX}trajectory-").removesuffix(".json") + end_key = f"{_PREFIX}{turn_id}/live/end.json" + assert sorted(run.store.list(_PREFIX)) == sorted([final_key, end_key]) + assert json.loads(run.store.get(run.store.object_url(end_key))) == { + "state": "completed", + "next": 0, + } + + @pytest.mark.asyncio async def test_without_live_trajectory_nothing_is_followed(monkeypatch, tmp_path): run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=False) From 07486830863a4342400ee4af5e9f97d0d792f686 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 8 Oct 2026 15:39:10 -0700 Subject: [PATCH 08/12] test(prompt_agent): follow a deployed agent's live trajectory over two turns The echo agent records its trajectory through the framework's log, with the same final bytes, and a live-steps prompt has it append steps a second apart. A two-turn conversation on the local defaults then checks each turn is stored under its own prefix while it runs and joins to its final trajectory. Co-Authored-By: Claude Opus 5.5 --- tst/data/a2a_agent/agent.py | 26 ++-- ...test_prompt_agent_live_trajectory_local.py | 128 ++++++++++++++++++ 2 files changed, 146 insertions(+), 8 deletions(-) create mode 100644 tst/integration/task_step/test_prompt_agent_live_trajectory_local.py diff --git a/tst/data/a2a_agent/agent.py b/tst/data/a2a_agent/agent.py index 35e2d253..e878863d 100644 --- a/tst/data/a2a_agent/agent.py +++ b/tst/data/a2a_agent/agent.py @@ -1,15 +1,19 @@ """The A2A agent the integration suites deploy when they need an agent but not a model. -Deterministic: it echoes the prompt and records a native trajectory. Agent-config, mcp-config, -trajectory and triggers come from the agentenv-protocol framework. ``model_params`` is a write-only -config field so the agent-config negotiation and its redaction on read-back can be observed. +Deterministic: it echoes the prompt and records the echo in the framework's trajectory log, which a +reader can follow while the turn runs; a ``live-steps `` in its prompt has it first append ``n`` step +events, a second apart. Agent-config, mcp-config, trajectory and triggers come from the agentenv-protocol +framework. ``model_params`` is a write-only config field so the agent-config negotiation and its +redaction on read-back can be observed. It also moves files both ways. After the echo it adds a ``read over : `` line for each file part it is sent, fetching only inline bytes and HTTP(S) URLs (``could not read ...`` otherwise), and it answers each ``send-file `` line of its prompt with a file part naming that URI. """ +import asyncio import base64 +import re from typing import Any from urllib.parse import urlsplit @@ -28,9 +32,13 @@ Usage, WriteOnly, a2a_agent, + enable, ) SEND_FILE = "send-file " +LIVE_STEPS = re.compile(r"live-steps (\d+)") +LIVE_STEP_SECONDS = 1 +TRAJECTORY_FORMAT = "agentenv-echo-agent/v1" READ_TIMEOUT_SECONDS = 120 @@ -64,24 +72,26 @@ async def _read(part: FilePart) -> str: version="1.0.0", ), config=EchoAgentConfig, - extensions=(MCP_CONFIG_V1, TRAJECTORY_V1, TRIGGERS_V1), + extensions=(MCP_CONFIG_V1, enable(TRAJECTORY_V1, live=True), TRIGGERS_V1), ) class EchoAgent(AgentEnvAgent): async def run(self, request: TaskRequest[EchoAgentConfig]) -> TaskResult: + request.trajectory.set_format(TRAJECTORY_FORMAT) prompt = "\n".join(part.text for part in request.parts if isinstance(part, TextPart)) + if steps := LIVE_STEPS.search(prompt): + for step in range(int(steps.group(1))): + request.trajectory.append({"type": "step", "step": step}) + await asyncio.sleep(LIVE_STEP_SECONDS) reply = f"Echo: {prompt}" reply = "\n".join([reply, *[await _read(part) for part in request.parts if isinstance(part, FilePart)]]) sends = [line.removeprefix(SEND_FILE).strip() for line in prompt.splitlines() if line.startswith(SEND_FILE)] files = [FilePart(uri=uri, name=uri.rsplit("/", 1)[-1], mime_type="text/plain") for uri in sends] + request.trajectory.append({"type": "echo", "input": prompt, "output": reply}) return ( TaskResult.builder() .succeeded() .parts([TextPart(text=reply), *files]) .usage(Usage(tool_call_count=0, input_tokens=len(prompt), output_tokens=len(reply))) - .native_trajectory( - format="agentenv-echo-agent/v1", - payload=[{"type": "echo", "input": prompt, "output": reply}], - ) .build() ) diff --git a/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py b/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py new file mode 100644 index 00000000..e3e3ae0b --- /dev/null +++ b/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py @@ -0,0 +1,128 @@ +"""prompt_agent follows a deployed agent's trajectory while each turn runs and stores it in chunks, on the local +defaults. The agents are the echo agent (tst/data/a2a_agent), which appends a step a second before its echo, so a +turn runs for a few seconds; over a two-turn conversation with a user-sim, each turn is stored under its own prefix, +in more than one chunk, and its chunks join to that turn's final trajectory. Needs Docker and a throwaway local +registry.""" + +import json +import re +import shutil +import socket +import subprocess +import time +import uuid + +import httpx +import pytest + +from agent_env.artifact.store import reset_artifact_store +from agent_env.config import configure, get_config, reset_config, set_image_store +from agent_env.store.image_store import LocalRegistryImageStore +from agent_env.store.object_store.local.grant_server import grant_server +from agent_env.task import Task +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep +from tst.util.a2a_test_agent import put_test_agent + +pytestmark = pytest.mark.integration + +REGISTRY_READY_SECONDS = 45 +# A turn runs for about this many seconds, a read a second apart, so the follower reads it more than once. +LIVE_STEPS = 4 + + +def _docker(*args): + return subprocess.run(["docker", *args], capture_output=True, text=True) + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@pytest.fixture +def local_registry(monkeypatch, tmp_path): + sandboxes = tmp_path / "sandboxes" + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_IMAGE_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) + port = _free_port() + name = f"live-trajectory-registry-{uuid.uuid4().hex[:8]}" + started = _docker("run", "-d", "--rm", "--name", name, "-p", f"127.0.0.1:{port}:5000", "registry:2") + assert started.returncode == 0, started.stderr + host = f"localhost:{port}" + deadline = time.time() + REGISTRY_READY_SECONDS + while time.time() < deadline: + try: + if httpx.get(f"http://{host}/v2/", timeout=2).status_code in (200, 401): + break + except httpx.HTTPError: + pass + time.sleep(0.5) + configure() + set_image_store(LocalRegistryImageStore(host)) + reset_artifact_store() + try: + yield tmp_path + finally: + # Only this test's containers: every local sandbox it deployed has a work dir in its own sandbox root, and its + # container is named after that sandbox. Images are shared, so matching by image would reach other runs. + for work_dir in sandboxes.glob("agent-env-*"): + if m := re.match(r"agent-env-(local-[0-9a-f]+)-", work_dir.name): + _docker("rm", "-f", f"agent-{m.group(1)}") + _docker("rm", "-f", name) + shutil.rmtree(sandboxes, ignore_errors=True) + # The agents uploaded through this process's grant server, whose certificate is from this test's state root; + # a later test's agents trust another root's CA, so it must start afresh. + grant_server(None, None).close() + reset_artifact_store() + reset_config() + + +def _deploy(agent, name: str) -> DeployAgentTaskStep: + return DeployAgentTaskStep( + id=f"deploy-{name}", version=None, env_ids=[], a2a_agent_id=agent.id, a2a_agent_version=agent.version, + agent_name=name, sandbox_type="local", + ) + + +def _json_lines(store, key: str) -> list: + return [json.loads(line) for line in store.get(store.object_url(key)).decode().splitlines()] + + +@pytest.mark.asyncio +async def test_each_turn_is_stored_live_under_its_own_prefix_and_joins_to_its_final_trajectory(local_registry): + suffix = uuid.uuid4().hex[:8] + store = get_config().get_object_store() + prefix_key = f"live-trajectory-{suffix}/" + agent = put_test_agent(f"live-trajectory-agent-{suffix}") + task = Task.put(id=f"live-trajectory-{suffix}", steps=[ + _deploy(agent, "solver"), + _deploy(agent, "user"), + PromptAgentTaskStep( + id="ask", version=None, prompt_id="ask", agent_name="solver", user_agent_name="user", + max_conversation_turns=2, poll_interval_seconds=1, prompt=f"live-steps {LIVE_STEPS}", + trajectory_output_prefix=store.object_url(prefix_key), live_trajectory=True, + ), + ]) + + context = await task.run() + + assert not context.metadata.get("failed_steps"), context.metadata.get("failed_steps") + finals = context.prompt_responses[-1].target_agent_per_turn_trajectory_s3_uris + assert len(finals) == 2 and all(finals) + turns = [store.get_object_key(url).removeprefix(f"{prefix_key}trajectory-").removesuffix(".json") for url in finals] + assert len(set(turns)) == 2 + for turn, final_url in zip(turns, finals): + final = json.loads(store.get(final_url)) + assert [event["type"] for event in final] == ["step"] * LIVE_STEPS + ["echo"] + live = f"{prefix_key}{turn}/live/" + chunks = sorted(key for key in store.list(live) if key.endswith(".jsonl")) + assert len(chunks) > 1, chunks + assert [event for key in chunks for event in _json_lines(store, key)] == final + assert json.loads(store.get(store.object_url(f"{live}meta.json"))) == {"format": "agentenv-echo-agent/v1"} + assert json.loads(store.get(store.object_url(f"{live}end.json"))) == {"state": "completed", "next": len(final)} From 3975417de77156575c1b6dba102fe37f873a0cb0 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 8 Oct 2026 15:40:30 -0700 Subject: [PATCH 09/12] fix(protocol): keep a trajectory's format once an event is appended Readers parse every event by the one format a read reports, so set_format now refuses a different format after the first append; naming the same format again is still allowed. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 1 + .../src/agentenv_protocol/a2a_agent/tasks/v1.py | 5 ++++- .../tests/test_a2a_trajectory_log.py | 12 ++++++++++++ 3 files changed, 17 insertions(+), 1 deletion(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 598f0da4..3f4108dd 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -327,6 +327,7 @@ never changes after it is appended and a read never returns part of one. It does no I/O; call it from one thread at a time. After the task ends the log refuses further events. `set_format` names the events' versioned format, which readers use to parse them; call it before the first `append`, which otherwise raises. +Once an event is appended the format cannot change. The framework keeps each task's log and serves it while the task runs. When the result carries no `native_trajectory`, the log is the task's final trajectory: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py index 2503f798..90fd4bc4 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/tasks/v1.py @@ -152,10 +152,13 @@ def format(self) -> str | None: return self._format def set_format(self, format: str) -> None: - """Name the events' versioned format, such as ``my-cli-events/1``.""" + """Name the events' versioned format, such as ``my-cli-events/1``; once an event is + appended, it stays.""" if not isinstance(format, str) or not format.strip(): raise ValueError("trajectory format must be a non-empty string") self._ensure_open() + if self._events and format != self._format: + raise RuntimeError("the trajectory's format cannot change once events are appended") self._format = format def append(self, event: Mapping[str, Any]) -> None: diff --git a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py index 84327bbb..175fc123 100644 --- a/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py +++ b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py @@ -425,6 +425,18 @@ def test_an_appended_event_never_changes() -> None: log.set_format(" ") +def test_the_format_cannot_change_once_an_event_is_appended() -> None: + log = TrajectoryLog() + log.set_format("first/1") + log.set_format("second/1") + log.append({"step": 1}) + + log.set_format("second/1") + with pytest.raises(RuntimeError, match="cannot change"): + log.set_format("third/1") + assert log.format == "second/1" + + def test_a_sealed_log_refuses_more_events() -> None: trajectories = TaskTrajectories() log = _open(trajectories, "task-1", "context-1") From abe739fdb391fc5c63e06d75bdea78419099eb75 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 8 Oct 2026 17:09:16 -0700 Subject: [PATCH 10/12] feat(prompt_agent): link a step's live turns in order Each turn's end.json ends that turn only, so a reader of a multi-turn step's stored live trajectory could not tell which turn came next or when the step was over. Each sent turn's next.json now names the turn sent after it, written as that turn's message goes out, and the step's last turn gets {"turn": null} once the step ends, on failure too. The turn loop is re-indented under the try that marks the last turn; the rest of its diff is whitespace only. Co-Authored-By: Claude Opus 5.5 --- .../snapshot_utils/live_trajectory.py | 121 ++++++-- .../task_step/task_steps/prompt_agent.py | 288 +++++++++--------- ...test_prompt_agent_live_trajectory_local.py | 6 +- .../a2a_agent/test_trajectory_follower.py | 50 +++ .../test_prompt_agent_live_trajectory.py | 20 +- 5 files changed, 311 insertions(+), 174 deletions(-) diff --git a/src/agent_env/task_step/snapshot_utils/live_trajectory.py b/src/agent_env/task_step/snapshot_utils/live_trajectory.py index 1d0e82dc..53519c23 100644 --- a/src/agent_env/task_step/snapshot_utils/live_trajectory.py +++ b/src/agent_env/task_step/snapshot_utils/live_trajectory.py @@ -4,6 +4,10 @@ ``canceled`` or ``failed`` when the step stopped first. A follower that stopped early leaves ``end.json``'s ``next`` short of the turn's events; the final trajectory holds them all. +A step's turns are linked in order: each sent turn's ``next.json`` is ``{"turn": ""}``, naming the +turn sent after it, or ``{"turn": null}`` once the step has ended, so a reader of one turn goes on to +the next and knows when the step is over. + Written from the worker with its own credentials, like the final trajectory. Every object is written once: one already there holds the same bytes, so a repeated write is a no-op.""" from __future__ import annotations @@ -11,7 +15,7 @@ import asyncio import json import logging -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable from contextlib import asynccontextmanager from agentenv_protocol.a2a_agent import TrajectoryState @@ -25,8 +29,8 @@ _WRITE_ATTEMPTS = 3 # The task has ended when the turn's wait returns, so draining its last events is normally one read. _DRAIN_SECONDS = 60 -# A stopped step marks its turn ended in passing: the write must not hold up the cancel. -_END_WRITE_SECONDS = 5 +# A step marks or links its turns in passing: the write must not hold up the turn or a cancel. +_PASSING_WRITE_SECONDS = 5 # The A2A task states a turn's trajectory ends in as themselves; any other (failed, rejected) ends it failed. _ENDED_AS = {"completed": TrajectoryState.COMPLETED, "canceled": TrajectoryState.CANCELED} @@ -63,18 +67,7 @@ async def end(self, state: TrajectoryState) -> None: self._ended = True async def _put(self, name: str, body: bytes, content_type: str) -> None: - key = self._directory + name - for attempt in range(1, _WRITE_ATTEMPTS + 1): - try: - await asyncio.to_thread(self._store.put, key, body, content_type=content_type) - return - except ObjectAlreadyExistsError: - return - except Exception as exc: - if attempt == _WRITE_ATTEMPTS: - raise - logger.warning("Live trajectory write of %s failed (attempt %d, will retry): %s", key, attempt, exc) - await asyncio.sleep(attempt) + await _put(self._store, self._directory + name, body, content_type) class LiveTurn: @@ -82,19 +75,32 @@ class LiveTurn: the state it ``ended`` in once the wait for it returns.""" def __init__( - self, endpoint: str, store: ObjectStore, prefix_key: str, turn_id: str, poll_interval_seconds: float + self, + endpoint: str, + store: ObjectStore, + prefix_key: str, + turn_id: str, + poll_interval_seconds: float, + previous_turn_id: str | None, ) -> None: self._endpoint = endpoint self._store = store self._prefix_key = prefix_key self._turn_id = turn_id self._poll_interval_seconds = poll_interval_seconds + self._previous_turn_id = previous_turn_id self.stop = asyncio.Event() self.followers: list[asyncio.Task] = [] self.sinks: list[LiveTrajectoryChunks] = [] + self.links: list[asyncio.Task] = [] self.ended_as: TrajectoryState | None = None def follow(self, task_id: str) -> None: + if self._previous_turn_id is not None: + linking = _link(self._store, self._prefix_key, self._previous_turn_id, self._turn_id) + self.links.append( + asyncio.create_task(_in_passing(linking, self._previous_turn_id, f"linked to turn {self._turn_id}")) + ) sink = LiveTrajectoryChunks(self._store, self._prefix_key, self._turn_id) self.sinks.append(sink) self.followers.append( @@ -118,15 +124,17 @@ async def following_live_trajectory( turn_id: str, *, poll_interval_seconds: float, + previous_turn_id: str | None = None, ) -> AsyncIterator[LiveTurn]: - """Yield the turn whose task is followed once its message is sent. + """Yield the turn whose task is followed once its message is sent; then ``previous_turn_id``'s + ``next.json`` names it. When the body returns, the follower drains what remains, for at most ``_DRAIN_SECONDS``; when it raises, the follower is stopped at once, since the agent may be gone. Either way, a turn whose end the follower did not read is marked ended in the state the body reported, else ``canceled`` for a cancelled step or ``failed``, so readers do not wait on it. A follower's own failure is logged and never fails the turn.""" - turn = LiveTurn(endpoint, store, prefix_key, turn_id, poll_interval_seconds) + turn = LiveTurn(endpoint, store, prefix_key, turn_id, poll_interval_seconds, previous_turn_id) stopped_as = TrajectoryState.FAILED try: @@ -154,17 +162,74 @@ async def following_live_trajectory( logger.warning( "Live trajectory of turn %s stopped (continuing): %r", turn_id, follower.exception() ) + if turn.links: + await asyncio.wait(turn.links) if turn.sinks: - await _mark_ended(turn.sinks[-1], turn.ended_as or stopped_as, turn_id) + state = turn.ended_as or stopped_as + await _in_passing(turn.sinks[-1].end(state), turn_id, f"marked {state.value}") + + +class LiveTurns: + """A step's turns, followed one after another: a turn sent after another is named in that one's + ``next.json``, and ``end`` marks the last turn sent as the step's last.""" + def __init__(self, endpoint: str, store: ObjectStore, prefix_key: str, *, poll_interval_seconds: float) -> None: + self._endpoint = endpoint + self._store = store + self._prefix_key = prefix_key + self._poll_interval_seconds = poll_interval_seconds + self._last_sent: str | None = None + + @asynccontextmanager + async def following(self, turn_id: str) -> AsyncIterator[LiveTurn]: + """``following_live_trajectory`` for the step's next turn.""" + async with following_live_trajectory( + self._endpoint, + self._store, + self._prefix_key, + turn_id, + poll_interval_seconds=self._poll_interval_seconds, + previous_turn_id=self._last_sent, + ) as turn: + try: + yield turn + finally: + if turn.sinks: + self._last_sent = turn_id -async def _mark_ended(sink: LiveTrajectoryChunks, state: TrajectoryState, turn_id: str) -> None: - marking = asyncio.create_task(sink.end(state)) - done, _ = await asyncio.wait({marking}, timeout=_END_WRITE_SECONDS) + async def end(self) -> None: + """Mark the last turn sent as the step's last, once the step has ended.""" + if self._last_sent is not None: + ending = _link(self._store, self._prefix_key, self._last_sent, None) + await _in_passing(ending, self._last_sent, "marked last") + + +async def _link(store: ObjectStore, prefix_key: str, turn_id: str, next_turn_id: str | None) -> None: + """Write ``turn_id``'s ``next.json``: the turn sent after it, or None after the step's last.""" + body = json.dumps({"turn": next_turn_id}).encode() + await _put(store, f"{prefix_key}{turn_id}/live/next.json", body, "application/json") + + +async def _put(store: ObjectStore, key: str, body: bytes, content_type: str) -> None: + for attempt in range(1, _WRITE_ATTEMPTS + 1): + try: + await asyncio.to_thread(store.put, key, body, content_type=content_type) + return + except ObjectAlreadyExistsError: + return + except Exception as exc: + if attempt == _WRITE_ATTEMPTS: + raise + logger.warning("Live trajectory write of %s failed (attempt %d, will retry): %s", key, attempt, exc) + await asyncio.sleep(attempt) + + +async def _in_passing(write: Awaitable[None], turn_id: str, what: str) -> None: + """Run ``write`` for at most ``_PASSING_WRITE_SECONDS``, logging rather than raising its failure.""" + writing = asyncio.ensure_future(write) + done, _ = await asyncio.wait({writing}, timeout=_PASSING_WRITE_SECONDS) if not done: - marking.cancel() - logger.warning("Live trajectory of turn %s was not marked %s in time (continuing)", turn_id, state.value) - elif marking.exception() is not None: - logger.warning( - "Live trajectory of turn %s was not marked %s (continuing): %r", turn_id, state.value, marking.exception() - ) + writing.cancel() + logger.warning("Live trajectory of turn %s was not %s in time (continuing)", turn_id, what) + elif writing.exception() is not None: + logger.warning("Live trajectory of turn %s was not %s (continuing): %r", turn_id, what, writing.exception()) diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index 882a8714..c3b726dc 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -39,7 +39,7 @@ trajectory_object_url, trajectory_prefix_key, ) -from agent_env.task_step.snapshot_utils.live_trajectory import following_live_trajectory +from agent_env.task_step.snapshot_utils.live_trajectory import LiveTurns from agent_env.task_step.snapshot_utils.snapshot_series import SnapshotConfig, SnapshotSeries logger = logging.getLogger(__name__) @@ -488,163 +488,165 @@ async def _execute_conversation( source_agent_per_turn_prompt_parts: list[list[dict] | None] = [] traj_ext_cached = A2AAgent.find_extension(card, A2AAgent.EXT_TRAJECTORY) live = self._live_trajectory_target(context, card, target_url, trajectory_output_prefix) + live_turns = ( + LiveTurns(*live, poll_interval_seconds=self.poll_interval_seconds) if live is not None else None + ) final_state: str = TaskState.completed.value trajectory_s3_uri: Optional[str] = None # What the target has been sent: of the files its replies name, the only ones a user-sim is made able # to read, so a reply naming any other object a store owns can't read it out through the user-sim. sent_to_target: set[str] = set() - for turn in range(self.max_conversation_turns): - # `target_a2a_task_id` is the client A2A message id sent to the target - # agent and recorded on the conversation as `a2a_task_id`. - # This id is also used as part of the key name for the trajectory S3 object. - target_a2a_task_id = uuid.uuid4().hex - - # The turn records the objects' own URLs, and only once the parts are ready to send, so an - # object the agent can't be sent leaves no turn waiting. - def record_turn() -> None: - sent_to_target.update(_file_uris(current_user_parts)) - conversation_store.add_a2a_task( + try: + for turn in range(self.max_conversation_turns): + # `target_a2a_task_id` is the client A2A message id sent to the target + # agent and recorded on the conversation as `a2a_task_id`. + # This id is also used as part of the key name for the trajectory S3 object. + target_a2a_task_id = uuid.uuid4().hex + + # The turn records the objects' own URLs, and only once the parts are ready to send, so an + # object the agent can't be sent leaves no turn waiting. + def record_turn() -> None: + sent_to_target.update(_file_uris(current_user_parts)) + conversation_store.add_a2a_task( + conversation_id=conversation_id, + parts=current_user_parts, + a2a_task_id=target_a2a_task_id, + role="user", + ) + # Turn 0 of a prompt-mode step re-sends exactly the text already + # persisted as PromptResponse.prompt_text; store a None placeholder + # instead of a second copy so the doc doesn't carry the prompt twice. + # The list stays index-aligned with the per-turn trajectory URIs. + source_agent_per_turn_prompt_parts.append( + None + if turn == 0 and _duplicates_prompt_text(current_user_parts, prompt_text) + else list(current_user_parts) + ) + + following = live_turns.following(target_a2a_task_id) if live_turns is not None else nullcontext(None) + async with following as live_turn: + sent_task_id, result = await send_and_wait( + target_url, current_user_parts, agent=agent, message_id=target_a2a_task_id, + context_id=solver_context_id, timeout_seconds=self.timeout_seconds, + poll_interval_seconds=self.poll_interval_seconds, before_send=record_turn, + on_sent=live_turn.follow if live_turn else None, + ) + target_state = result["status"]["state"] + if live_turn: + live_turn.ended(target_state) + status_msg = (result.get("status") or {}).get("message") or {} + final_terminal = protocol.TerminalResponse.from_message(status_msg) + agent_response_parts = status_msg.get("parts") or [{"kind": "text", "text": final_terminal.response_text}] + + conversation_store.complete_a2a_task( conversation_id=conversation_id, - parts=current_user_parts, - a2a_task_id=target_a2a_task_id, - role="user", - ) - # Turn 0 of a prompt-mode step re-sends exactly the text already - # persisted as PromptResponse.prompt_text; store a None placeholder - # instead of a second copy so the doc doesn't carry the prompt twice. - # The list stays index-aligned with the per-turn trajectory URIs. - source_agent_per_turn_prompt_parts.append( - None - if turn == 0 and _duplicates_prompt_text(current_user_parts, prompt_text) - else list(current_user_parts) + parts=agent_response_parts, + role="agent", ) - following = ( - following_live_trajectory( - *live, target_a2a_task_id, poll_interval_seconds=self.poll_interval_seconds - ) - if live is not None - else nullcontext(None) - ) - async with following as live_turn: - sent_task_id, result = await send_and_wait( - target_url, current_user_parts, agent=agent, message_id=target_a2a_task_id, - context_id=solver_context_id, timeout_seconds=self.timeout_seconds, - poll_interval_seconds=self.poll_interval_seconds, before_send=record_turn, - on_sent=live_turn.follow if live_turn else None, - ) - target_state = result["status"]["state"] - if live_turn: - live_turn.ended(target_state) - status_msg = (result.get("status") or {}).get("message") or {} - final_terminal = protocol.TerminalResponse.from_message(status_msg) - agent_response_parts = status_msg.get("parts") or [{"kind": "text", "text": final_terminal.response_text}] - - conversation_store.complete_a2a_task( - conversation_id=conversation_id, - parts=agent_response_parts, - role="agent", - ) + if traj_ext_cached: + turn_traj_uri: str | None = None + try: + turn_traj_uri = await self._fetch_trajectory( + target_url, traj_ext_cached, sent_task_id, trajectory_output_prefix, + target_a2a_task_id, card=card, sandbox_type=agent.sandbox_type, + ) + except Exception as e: + logger.warning(f"Turn {turn+1} trajectory fetch failed (continuing): {e}") + per_turn_trajectory_s3_uris.append(turn_traj_uri) + + final_task_id = sent_task_id + final_state = target_state + logger.info(f"Turn {turn+1}/{self.max_conversation_turns} agent response ({target_state}): {final_terminal.response_text[:120]}...") + + if target_state == TaskState.failed: + conversation_store.mark_closed(conversation_id) + break + if turn == self.max_conversation_turns - 1: + break + + conv_doc = conversation_store.get_conversation(conversation_id) + if conv_doc and conv_doc.get("status") != "active": + logger.info(f"Conversation {conversation_id} no longer active (status={conv_doc.get('status')}); ending multi-turn") + break - if traj_ext_cached: - turn_traj_uri: str | None = None + # After every `break` above: a turn that ends the conversation is covered + # by teardown's final capture instead of being captured twice. + if series is not None: + await series.capture_turn(context, turn + 1) + + if trig_ext and triggers_registered: + env_triggers = await self._read_env_triggers(context) + (context.metadata.setdefault("env_trigger_snapshots", {}) + .setdefault(self.id, []).append({"turn": turn + 1, "envs": env_triggers})) + resp = await self._decide(decide_url, turn + 1, final_terminal.response_text, conversation_id, env_triggers) + (context.metadata.setdefault("agent_trigger_firings", {}) + .setdefault(self.id, []).append({"turn": turn + 1, "fired": resp.get("fired", [])})) + logger.info(f"Turn {turn+1} agent-trigger decide: fired={resp.get('fired', [])} done={resp.get('done')}") + if bool(resp.get("done")): + current_user_parts = resp.get("parts") or [] + conversation_store.add_a2a_task(conversation_id, parts=current_user_parts, + a2a_task_id=uuid.uuid4().hex, role="user") + conversation_store.mark_closed(conversation_id) + break + if resp.get("fired"): + current_user_parts = resp.get("parts") or [] + continue + # no trigger fired -> fall through to the LLM user-sim for this turn (hybrid) + + user_a2a_task_id = uuid.uuid4().hex try: - turn_traj_uri = await self._fetch_trajectory( - target_url, traj_ext_cached, sent_task_id, trajectory_output_prefix, - target_a2a_task_id, card=card, sandbox_type=agent.sandbox_type, + # A user-sim runs in a sandbox agent-env deployed; a human peer, registered or named by + # user_a2a_url, has none and reads the store itself. + _, user_result = await send_and_wait( + user_url, agent_response_parts, + agent=user_sim if is_user_sim and user_sim.sandbox_id else None, shareable=sent_to_target, + message_id=user_a2a_task_id, context_id=conversation_id, + timeout_seconds=self.user_agent_timeout_seconds, + poll_interval_seconds=self.poll_interval_seconds, ) - except Exception as e: - logger.warning(f"Turn {turn+1} trajectory fetch failed (continuing): {e}") - per_turn_trajectory_s3_uris.append(turn_traj_uri) - - final_task_id = sent_task_id - final_state = target_state - logger.info(f"Turn {turn+1}/{self.max_conversation_turns} agent response ({target_state}): {final_terminal.response_text[:120]}...") - - if target_state == TaskState.failed: - conversation_store.mark_closed(conversation_id) - break - if turn == self.max_conversation_turns - 1: - break - - conv_doc = conversation_store.get_conversation(conversation_id) - if conv_doc and conv_doc.get("status") != "active": - logger.info(f"Conversation {conversation_id} no longer active (status={conv_doc.get('status')}); ending multi-turn") - break - - # After every `break` above: a turn that ends the conversation is covered - # by teardown's final capture instead of being captured twice. - if series is not None: - await series.capture_turn(context, turn + 1) - - if trig_ext and triggers_registered: - env_triggers = await self._read_env_triggers(context) - (context.metadata.setdefault("env_trigger_snapshots", {}) - .setdefault(self.id, []).append({"turn": turn + 1, "envs": env_triggers})) - resp = await self._decide(decide_url, turn + 1, final_terminal.response_text, conversation_id, env_triggers) - (context.metadata.setdefault("agent_trigger_firings", {}) - .setdefault(self.id, []).append({"turn": turn + 1, "fired": resp.get("fired", [])})) - logger.info(f"Turn {turn+1} agent-trigger decide: fired={resp.get('fired', [])} done={resp.get('done')}") - if bool(resp.get("done")): - current_user_parts = resp.get("parts") or [] - conversation_store.add_a2a_task(conversation_id, parts=current_user_parts, - a2a_task_id=uuid.uuid4().hex, role="user") + except TimeoutError as e: + logger.warning(f"user_a2a_url timeout for conversation {conversation_id} ({e}); marking abandoned") conversation_store.mark_closed(conversation_id) break - if resp.get("fired"): - current_user_parts = resp.get("parts") or [] - continue - # no trigger fired -> fall through to the LLM user-sim for this turn (hybrid) - user_a2a_task_id = uuid.uuid4().hex - try: - # A user-sim runs in a sandbox agent-env deployed; a human peer, registered or named by - # user_a2a_url, has none and reads the store itself. - _, user_result = await send_and_wait( - user_url, agent_response_parts, - agent=user_sim if is_user_sim and user_sim.sandbox_id else None, shareable=sent_to_target, - message_id=user_a2a_task_id, context_id=conversation_id, - timeout_seconds=self.user_agent_timeout_seconds, - poll_interval_seconds=self.poll_interval_seconds, - ) - except TimeoutError as e: - logger.warning(f"user_a2a_url timeout for conversation {conversation_id} ({e}); marking abandoned") - conversation_store.mark_closed(conversation_id) - break - - user_state = (user_result.get("status") or {}).get("state") - if user_state != TaskState.completed: - logger.info(f"user_a2a_url returned state={user_state}; ending multi-turn") - conversation_store.mark_closed(conversation_id) - break - - user_msg = (user_result.get("status") or {}).get("message") or {} - current_user_parts = user_msg.get("parts") or [] - if not current_user_parts: - logger.warning(f"user_a2a_url returned no parts; ending multi-turn") - conversation_store.mark_closed(conversation_id) - break - - if is_user_sim: - raw_text = next((p["text"] for p in current_user_parts if p.get("kind") == "text"), "") - try: - parsed = json.loads(raw_text) - message_text = parsed["message"] - done = bool(parsed["done"]) - except (json.JSONDecodeError, KeyError) as e: - logger.warning(f"user-sim returned malformed output despite output_format ({e}); using raw text as next prompt; raw={raw_text[:200]!r}") - current_user_parts = [{"kind": "text", "text": raw_text}] - else: - extras = {k: v for k, v in parsed.items() if k not in ("message", "done")} - if extras: - (context.metadata.setdefault("usersim_turn_outputs", {}) - .setdefault(self.id, []).append({"turn": turn + 1, "fields": extras})) - current_user_parts = [{"kind": "text", "text": message_text}] - if done: - logger.info(f"user-sim signaled done=true at turn {turn+1}; ending multi-turn") - conversation_store.mark_closed(conversation_id) - break + user_state = (user_result.get("status") or {}).get("state") + if user_state != TaskState.completed: + logger.info(f"user_a2a_url returned state={user_state}; ending multi-turn") + conversation_store.mark_closed(conversation_id) + break + + user_msg = (user_result.get("status") or {}).get("message") or {} + current_user_parts = user_msg.get("parts") or [] + if not current_user_parts: + logger.warning(f"user_a2a_url returned no parts; ending multi-turn") + conversation_store.mark_closed(conversation_id) + break + + if is_user_sim: + raw_text = next((p["text"] for p in current_user_parts if p.get("kind") == "text"), "") + try: + parsed = json.loads(raw_text) + message_text = parsed["message"] + done = bool(parsed["done"]) + except (json.JSONDecodeError, KeyError) as e: + logger.warning(f"user-sim returned malformed output despite output_format ({e}); using raw text as next prompt; raw={raw_text[:200]!r}") + current_user_parts = [{"kind": "text", "text": raw_text}] + else: + extras = {k: v for k, v in parsed.items() if k not in ("message", "done")} + if extras: + (context.metadata.setdefault("usersim_turn_outputs", {}) + .setdefault(self.id, []).append({"turn": turn + 1, "fields": extras})) + current_user_parts = [{"kind": "text", "text": message_text}] + if done: + logger.info(f"user-sim signaled done=true at turn {turn+1}; ending multi-turn") + conversation_store.mark_closed(conversation_id) + break + finally: + # On failure too, so a reader of the last turn knows no other follows. + if live_turns is not None: + await live_turns.end() await self._persist_env_trigger_state(context) diff --git a/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py b/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py index e3e3ae0b..c81601a3 100644 --- a/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py +++ b/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py @@ -1,8 +1,8 @@ """prompt_agent follows a deployed agent's trajectory while each turn runs and stores it in chunks, on the local defaults. The agents are the echo agent (tst/data/a2a_agent), which appends a step a second before its echo, so a turn runs for a few seconds; over a two-turn conversation with a user-sim, each turn is stored under its own prefix, -in more than one chunk, and its chunks join to that turn's final trajectory. Needs Docker and a throwaway local -registry.""" +in more than one chunk, and its chunks join to that turn's final trajectory; each turn names the one after it, and +the last none. Needs Docker and a throwaway local registry.""" import json import re @@ -126,3 +126,5 @@ async def test_each_turn_is_stored_live_under_its_own_prefix_and_joins_to_its_fi assert [event for key in chunks for event in _json_lines(store, key)] == final assert json.loads(store.get(store.object_url(f"{live}meta.json"))) == {"format": "agentenv-echo-agent/v1"} assert json.loads(store.get(store.object_url(f"{live}end.json"))) == {"state": "completed", "next": len(final)} + links = [json.loads(store.get(store.object_url(f"{prefix_key}{turn}/live/next.json"))) for turn in turns] + assert links == [{"turn": turns[1]}, {"turn": None}] diff --git a/tst/unit/a2a_agent/test_trajectory_follower.py b/tst/unit/a2a_agent/test_trajectory_follower.py index 02627a6e..fdeae5f3 100644 --- a/tst/unit/a2a_agent/test_trajectory_follower.py +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -15,6 +15,7 @@ from agent_env.task_step.snapshot_utils import live_trajectory from agent_env.task_step.snapshot_utils.live_trajectory import ( LiveTrajectoryChunks, + LiveTurns, following_live_trajectory, ) @@ -303,3 +304,52 @@ async def test_a_reported_task_state_ends_the_turn_as_its_trajectory_state(monke turn.ended(task_state) assert _end(store) == {"state": state, "next": 0} + + +def _next(store: LocalFilesystemObjectStore, turn: str) -> dict: + return json.loads(store.get(store.object_url(f"prefix/{turn}/live/next.json"))) + + +@pytest.mark.asyncio +async def test_each_turn_names_the_next_and_the_last_is_marked_last(monkeypatch, tmp_path): + _agent_answering(monkeypatch, _page(0, [{"i": 0}], "completed")) + store = LocalFilesystemObjectStore(str(tmp_path)) + turns = LiveTurns(_ENDPOINT, store, "prefix/", poll_interval_seconds=0) + + for turn_id in ("turn-1", "turn-2"): + async with turns.following(turn_id) as turn: + turn.follow(f"task-{turn_id}") + await _until_stored(store, f"prefix/{turn_id}/live/end.json") + turn.ended("completed") + assert "prefix/turn-2/live/next.json" not in store.list("prefix/") + await turns.end() + + assert _next(store, "turn-1") == {"turn": "turn-2"} + assert _next(store, "turn-2") == {"turn": None} + + +@pytest.mark.asyncio +async def test_a_turn_never_sent_is_not_linked_and_the_last_sent_is_marked_last(monkeypatch, tmp_path): + _agent_answering(monkeypatch, _page(0, [{"i": 0}], "completed")) + store = LocalFilesystemObjectStore(str(tmp_path)) + turns = LiveTurns(_ENDPOINT, store, "prefix/", poll_interval_seconds=0) + async with turns.following("turn-1") as turn: + turn.follow("task-1") + turn.ended("completed") + + with pytest.raises(RuntimeError): + async with turns.following("turn-2"): + raise RuntimeError("the message could not be sent") + await turns.end() + + assert _next(store, "turn-1") == {"turn": None} + assert not any(key.startswith("prefix/turn-2/") for key in store.list("prefix/")) + + +@pytest.mark.asyncio +async def test_a_step_that_sent_no_turn_marks_none(tmp_path): + store = LocalFilesystemObjectStore(str(tmp_path)) + + await LiveTurns(_ENDPOINT, store, "prefix/", poll_interval_seconds=0).end() + + assert store.list("prefix/") == [] diff --git a/tst/unit/task_step/test_prompt_agent_live_trajectory.py b/tst/unit/task_step/test_prompt_agent_live_trajectory.py index 84111bad..7b34b48f 100644 --- a/tst/unit/task_step/test_prompt_agent_live_trajectory.py +++ b/tst/unit/task_step/test_prompt_agent_live_trajectory.py @@ -117,6 +117,7 @@ async def test_a_live_turn_is_stored_in_chunks_while_it_runs(monkeypatch, tmp_pa "state": "completed", "next": len(_EVENTS + _MORE), } + assert json.loads(run.store.get(run.store.object_url(f"{live}next.json"))) == {"turn": None} @pytest.mark.asyncio @@ -132,13 +133,30 @@ async def stops_at_once(*_args, **_kwargs) -> None: final_key = run.store.get_object_key(result.prompt_responses[-1].agent_trajectory_s3_uri) turn_id = final_key.removeprefix(f"{_PREFIX}trajectory-").removesuffix(".json") end_key = f"{_PREFIX}{turn_id}/live/end.json" - assert sorted(run.store.list(_PREFIX)) == sorted([final_key, end_key]) + next_key = f"{_PREFIX}{turn_id}/live/next.json" + assert sorted(run.store.list(_PREFIX)) == sorted([final_key, end_key, next_key]) assert json.loads(run.store.get(run.store.object_url(end_key))) == { "state": "completed", "next": 0, } +@pytest.mark.asyncio +async def test_a_step_that_fails_after_a_turn_marks_that_turn_last(monkeypatch, tmp_path): + run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=False) + + def unavailable(*_args, **_kwargs) -> None: + raise RuntimeError("conversation store unavailable") + + monkeypatch.setattr(pa.conversation_store, "complete_a2a_task", unavailable) + + with pytest.raises(RuntimeError, match="conversation store unavailable"): + await run.execute(run.step(live_trajectory=True)) + + (next_key,) = [key for key in run.store.list(_PREFIX) if key.endswith("/live/next.json")] + assert json.loads(run.store.get(run.store.object_url(next_key))) == {"turn": None} + + @pytest.mark.asyncio async def test_without_live_trajectory_nothing_is_followed(monkeypatch, tmp_path): run = _Run(monkeypatch, tmp_path, live=True, wait_for_chunk=False) From c90ba25fc7501dca7563b27faaa31ca9f748bcb8 Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 8 Oct 2026 17:32:51 -0700 Subject: [PATCH 11/12] test(prompt_agent): read the live trajectory tests' results by their object-url names main renamed PromptResponse's trajectory fields to object-url names, so the live trajectory tests read agent_trajectory_object_url and target_agent_per_turn_trajectory_object_urls. Co-Authored-By: Claude Opus 5.5 --- .../task_step/test_prompt_agent_live_trajectory_local.py | 2 +- tst/unit/task_step/test_prompt_agent_live_trajectory.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py b/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py index e3e3ae0b..07c05c87 100644 --- a/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py +++ b/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py @@ -113,7 +113,7 @@ async def test_each_turn_is_stored_live_under_its_own_prefix_and_joins_to_its_fi context = await task.run() assert not context.metadata.get("failed_steps"), context.metadata.get("failed_steps") - finals = context.prompt_responses[-1].target_agent_per_turn_trajectory_s3_uris + finals = context.prompt_responses[-1].target_agent_per_turn_trajectory_object_urls assert len(finals) == 2 and all(finals) turns = [store.get_object_key(url).removeprefix(f"{prefix_key}trajectory-").removesuffix(".json") for url in finals] assert len(set(turns)) == 2 diff --git a/tst/unit/task_step/test_prompt_agent_live_trajectory.py b/tst/unit/task_step/test_prompt_agent_live_trajectory.py index 84111bad..caaad2a1 100644 --- a/tst/unit/task_step/test_prompt_agent_live_trajectory.py +++ b/tst/unit/task_step/test_prompt_agent_live_trajectory.py @@ -100,7 +100,7 @@ async def test_a_live_turn_is_stored_in_chunks_while_it_runs(monkeypatch, tmp_pa result = await run.execute(run.step(live_trajectory=True)) - final_url = result.prompt_responses[-1].agent_trajectory_s3_uri + final_url = result.prompt_responses[-1].agent_trajectory_object_url turn_id = run.store.get_object_key(final_url).removeprefix(f"{_PREFIX}trajectory-").removesuffix(".json") live = f"{_PREFIX}{turn_id}/live/" chunks = _chunk_keys(run.store) @@ -129,7 +129,7 @@ async def stops_at_once(*_args, **_kwargs) -> None: result = await run.execute(run.step(live_trajectory=True)) - final_key = run.store.get_object_key(result.prompt_responses[-1].agent_trajectory_s3_uri) + final_key = run.store.get_object_key(result.prompt_responses[-1].agent_trajectory_object_url) turn_id = final_key.removeprefix(f"{_PREFIX}trajectory-").removesuffix(".json") end_key = f"{_PREFIX}{turn_id}/live/end.json" assert sorted(run.store.list(_PREFIX)) == sorted([final_key, end_key]) @@ -147,7 +147,7 @@ async def test_without_live_trajectory_nothing_is_followed(monkeypatch, tmp_path assert run.cursor_reads == [] assert _chunk_keys(run.store) == [] - assert json.loads(run.store.get(result.prompt_responses[-1].agent_trajectory_s3_uri)) == _EVENTS + _MORE + assert json.loads(run.store.get(result.prompt_responses[-1].agent_trajectory_object_url)) == _EVENTS + _MORE @pytest.mark.asyncio @@ -169,7 +169,7 @@ async def test_an_agent_without_the_live_read_is_not_followed(monkeypatch, tmp_p assert run.cursor_reads == [] assert _chunk_keys(run.store) == [] - assert json.loads(run.store.get(result.prompt_responses[-1].agent_trajectory_s3_uri)) == _EVENTS + _MORE + assert json.loads(run.store.get(result.prompt_responses[-1].agent_trajectory_object_url)) == _EVENTS + _MORE def test_a_prefix_outside_the_store_is_not_followed(monkeypatch, tmp_path, caplog): From b19695877df55c249369840307ffb52b32e699ca Mon Sep 17 00:00:00 2001 From: Can Ersoz Date: Thu, 8 Oct 2026 17:59:13 -0700 Subject: [PATCH 12/12] test(a2a_agent): bound the echo agent's live steps A prompt that happened to contain a large live-steps count held its turn for that many seconds; the echo agent now appends at most ten. Co-Authored-By: Claude Opus 5.5 --- tst/data/a2a_agent/agent.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tst/data/a2a_agent/agent.py b/tst/data/a2a_agent/agent.py index e878863d..b1a12f6a 100644 --- a/tst/data/a2a_agent/agent.py +++ b/tst/data/a2a_agent/agent.py @@ -2,9 +2,9 @@ Deterministic: it echoes the prompt and records the echo in the framework's trajectory log, which a reader can follow while the turn runs; a ``live-steps `` in its prompt has it first append ``n`` step -events, a second apart. Agent-config, mcp-config, trajectory and triggers come from the agentenv-protocol -framework. ``model_params`` is a write-only config field so the agent-config negotiation and its -redaction on read-back can be observed. +events (at most ``MAX_LIVE_STEPS``), a second apart. Agent-config, mcp-config, trajectory and triggers +come from the agentenv-protocol framework. ``model_params`` is a write-only config field so the +agent-config negotiation and its redaction on read-back can be observed. It also moves files both ways. After the echo it adds a ``read over : `` line for each file part it is sent, fetching only inline bytes and HTTP(S) URLs (``could not read ...`` otherwise), @@ -38,6 +38,8 @@ SEND_FILE = "send-file " LIVE_STEPS = re.compile(r"live-steps (\d+)") LIVE_STEP_SECONDS = 1 +# Bounds how long a prompt can hold up a turn. +MAX_LIVE_STEPS = 10 TRAJECTORY_FORMAT = "agentenv-echo-agent/v1" READ_TIMEOUT_SECONDS = 120 @@ -79,7 +81,7 @@ async def run(self, request: TaskRequest[EchoAgentConfig]) -> TaskResult: request.trajectory.set_format(TRAJECTORY_FORMAT) prompt = "\n".join(part.text for part in request.parts if isinstance(part, TextPart)) if steps := LIVE_STEPS.search(prompt): - for step in range(int(steps.group(1))): + for step in range(min(int(steps.group(1)), MAX_LIVE_STEPS)): request.trajectory.append({"type": "step", "step": step}) await asyncio.sleep(LIVE_STEP_SECONDS) reply = f"Echo: {prompt}"