diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index c11b4d35..b2e7fb66 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,77 @@ 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; 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: +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 + `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: + 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, 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 + 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 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 @@ -319,7 +390,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 +488,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 +628,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..0240712f --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/_trajectory_log.py @@ -0,0 +1,244 @@ +"""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_CONTEXTS = 1_024 +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 + + +@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. + + 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_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() + + 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._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.""" + 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=next( + (record.log.format for record in reversed(records) if record.log.format), + None, + ), + 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, *, keep: str) -> None: + for context_id in tuple(self._contexts): + if len(self._contexts) <= self._max_contexts: + return + task_ids = self._contexts[context_id] + 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: + del self._tasks[task_id] + + +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..a0a7db36 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,45 @@ 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" + + +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, le=MAX_TRAJECTORY_READ_EVENTS) + + +class ContextEventsTrajectoryRequest(ExtensionRequest): + context_id: str + after: StrictInt = Field(ge=0) + limit: StrictInt | None = Field(default=None, ge=1, le=MAX_TRAJECTORY_READ_EVENTS) + + +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 +613,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 +621,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 +920,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 +1032,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 +1231,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..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,20 +28,28 @@ from starlette.routing import Route from starlette.types import ASGIApp, Receive, Scope, Send +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, TRIGGERS_V1, + ContextEventsTrajectoryRequest, + ContextObjectTrajectoryRequest, + ContextTrajectoryRequest, ExtensionActivation, ExtensionDefinition, ImplementationOwner, McpAddRequest, OperationReference, + TaskEventsTrajectoryRequest, TaskObjectTrajectoryRequest, TaskTrajectoryRequest, + TrajectoryEventsResponse, + TrajectoryState, TriggerDecideRequest, TriggerRegisterRequest, enable, @@ -89,7 +97,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 +162,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 +433,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 +538,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 +728,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 +764,38 @@ async def trajectory_get( } } + def _trajectory_events( + self, request: TaskEventsTrajectoryRequest | ContextEventsTrajectoryRequest + ) -> TrajectoryEventsResponse: + 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 + ) + 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 +1149,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 +1163,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 +1176,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 +1188,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 +1294,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..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 @@ -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,67 @@ 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``; 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: + """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) + 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 +201,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..175fc123 --- /dev/null +++ b/packages/agentenv-protocol/tests/test_a2a_trajectory_log.py @@ -0,0 +1,634 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import threading +import time +from datetime import datetime, timedelta, timezone +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, + 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 + +_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 _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)] + + +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}, + {"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: + 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.set_format("test-events/1") + 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 = _open(trajectories, "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.set_format("test-events/1") + 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.set_format("test-events/1") + 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.set_format("test-events/1") + 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 = _open(services.task_trajectories, "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() + log.set_format("test-events/1") + 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_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") + 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 = _open(trajectories, "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 = _open(trajectories, "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 = _open(trajectories, "queued", "context-1") + running = _open(trajectories, "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_eviction_drops_completed_contexts_but_never_running_ones() -> None: + trajectories = TaskTrajectories(max_contexts=2) + running = _open(trajectories, "running", "context-running") + trajectories.start("running") + 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"[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 = _open(trajectories, "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( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(tasks_v1, "_LARGE_LIVE_TRAJECTORY_BYTES", 10) + 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", + ) + + +# 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 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..53519c23 --- /dev/null +++ b/src/agent_env/task_step/snapshot_utils/live_trajectory.py @@ -0,0 +1,235 @@ +"""A turn's live trajectory, stored as it is read: ``/live/-.jsonl`` per +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. + +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 + +import asyncio +import json +import logging +from collections.abc import AsyncIterator, Awaitable +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 + +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 +# 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} + + +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 + self._stored = 0 + self._ended = 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") + 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: + 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: + await _put(self._store, self._directory + name, body, content_type) + + +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, + 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( + 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, + store: ObjectStore, + prefix_key: str, + 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; 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, previous_turn_id) + stopped_as = TrajectoryState.FAILED + + try: + 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)", + turn_id, + _DRAIN_SECONDS, + ) + except asyncio.CancelledError: + stopped_as = TrajectoryState.CANCELED + raise + finally: + turn.stop.set() + for follower in turn.followers: + follower.cancel() + 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 turn.links: + await asyncio.wait(turn.links) + if turn.sinks: + 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 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: + 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 3765cf4a..3528a0df 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 LiveTurns 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,152 +487,166 @@ async def _execute_conversation( per_turn_trajectory_urls: 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) + 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_url: 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", ) - 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, - ) - target_state = result["status"]["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}] + 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_urls.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 - conversation_store.complete_a2a_task( - conversation_id=conversation_id, - parts=agent_response_parts, - role="agent", - ) + 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_urls.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) @@ -719,6 +743,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/data/a2a_agent/agent.py b/tst/data/a2a_agent/agent.py index 35e2d253..b1a12f6a 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 (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), 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,15 @@ Usage, WriteOnly, a2a_agent, + enable, ) 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 @@ -64,24 +74,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(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}" 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..b82bd60f --- /dev/null +++ b/tst/integration/task_step/test_prompt_agent_live_trajectory_local.py @@ -0,0 +1,130 @@ +"""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; each turn names the one after it, and +the last none. 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_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 + 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)} + 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 new file mode 100644 index 00000000..fdeae5f3 --- /dev/null +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -0,0 +1,355 @@ +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, + LiveTurns, + 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 + + +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: + 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, + "AsyncClient", + 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() + + with pytest.raises(type(error)): + async with following_live_trajectory( + _ENDPOINT, store, "prefix/", "turn", poll_interval_seconds=0 + ) as turn: + turn.follow("task-1") + await asyncio.sleep(0.05) + raise error + + assert time.monotonic() - started < 2 + 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 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) + 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 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} + + +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 new file mode 100644 index 00000000..99e04ff6 --- /dev/null +++ b/tst/unit/task_step/test_prompt_agent_live_trajectory.py @@ -0,0 +1,213 @@ +"""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.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 + +_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_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) + 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), + } + assert json.loads(run.store.get(run.store.object_url(f"{live}next.json"))) == {"turn": None} + + +@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_object_url) + turn_id = final_key.removeprefix(f"{_PREFIX}trajectory-").removesuffix(".json") + end_key = f"{_PREFIX}{turn_id}/live/end.json" + 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) + + 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_object_url)) == _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_object_url)) == _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