diff --git a/agent_core/components/observers/trajectory.py b/agent_core/components/observers/trajectory.py index 7ab42c2..8cd3981 100644 --- a/agent_core/components/observers/trajectory.py +++ b/agent_core/components/observers/trajectory.py @@ -11,7 +11,8 @@ import json import os import time -from collections.abc import Callable, Iterable +from collections import deque +from collections.abc import Callable, Iterable, Mapping from pathlib import Path from typing import IO, Any, Literal @@ -24,6 +25,7 @@ ToolResult, TurnContext, ) +from agent_core.messages import Message from agent_core.tool_content import redacted_for_trace, redacted_tool_result_content _FORMATS: tuple[str, ...] = ("json", "jsonl") @@ -31,7 +33,38 @@ # Normal completion markers (Anthropic / OpenAI-compatible). They carry no # information and appear on almost every turn, so trajectories omit them. _ORDINARY_FINISH_REASONS = frozenset({"end_turn", "stop"}) -_STREAM_ENCODER = json.JSONEncoder(ensure_ascii=False, separators=(",", ":")) +_DIAGNOSTIC_KEYS = ("usage", "thinking_blocks", "finish_reason", "stop_details") +# Blocks a signature-recovery reset strips from assistant content. +_THINKING_BLOCK_TYPES = frozenset({"thinking", "redacted_thinking"}) + + +def _json_default(value: Any) -> Any: + """Degrade a non-JSON value instead of failing every later snapshot. + + Loop history and tool results may carry host metadata (datetimes, SDK + objects, bytes). A trace is diagnostic, so a lossy rendering beats losing + the whole document to one ``TypeError``. + """ + if isinstance(value, Mapping): + return dict(value) + if isinstance(value, (bytes, bytearray)): + return f"<{len(value)} bytes>" + if isinstance(value, (set, frozenset)): + return list(value) + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + with contextlib.suppress(Exception): + return model_dump(mode="json") + isoformat = getattr(value, "isoformat", None) + if callable(isoformat): + with contextlib.suppress(Exception): + return isoformat() + return str(value) + + +_STREAM_ENCODER = json.JSONEncoder( + ensure_ascii=False, separators=(",", ":"), default=_json_default, +) _DEFAULT_FORMAT_ENV_VARS = ( "AGENT_CORE_TRAJECTORY_FORMATS", "FRONTIER_AGENT_TRAJECTORY_FORMATS", @@ -54,7 +87,7 @@ def _env_int(name: str, default: int) -> int: # ── Trajectory memory guardrails ────── # # One instance of this observer exists per agent (main + every sub-agent). -# Messages are spooled one-per-line to disk, and two additional bounds keep +# History is streamed one message at a time to disk, and two bounds keep # each public JSON snapshot cheap: # # _BODY_MAX_CHARS — per tool-result body kept in the JSON envelope. @@ -63,9 +96,11 @@ def _env_int(name: str, default: int) -> int: # JSONL is enabled (the default), that stream still records every body # untruncated. JSON-only configurations intentionally keep only the cap. # _COALESCE_N / _COALESCE_MS — snapshot batching, mirroring the -# equivalent defaults in the external worker-trace observer. -# Bounds the O(n²) whole-document copy to -# one per batch. The copy is streamed without a second full string. +# equivalent defaults in the external worker-trace observer. Every +# history-changing event (LLM input/response, tool result, turn end) +# counts once. Bounds the O(n²) whole-document copy to one per batch. +# The copy is streamed without a second full string, and only messages +# new or changed since the previous snapshot are re-rendered. # # 0 disables the respective bound. _BODY_MAX_CHARS: int = _env_int("SWARM_TRAJECTORY_BODY_MAX_CHARS", 16384) @@ -101,9 +136,20 @@ def _resolve_formats( class TrajectoryFileObserver(BaseObserver): - """Saves an agent's main / sub-agent traces in one or more formats.""" + """Save conversation snapshots (JSON) and execution events (JSONL). + + JSON follows the loop's current history, including rewrites and synthetic + replies, with diagnostic annotations, image redaction and the configured + tool-body cap. JSONL retains original events even after history is pruned. + """ - critical: bool = False + # These hooks do synchronous file I/O without yielding, so a background + # task never moved that work off the event loop; it only reordered it. + # Awaiting them preserves event order even when a synthetic result needs + # no async tool execution, where a background on_llm_response could run + # AFTER it. Coalescing and re-rendering only new or changed messages keep + # the awaited cost per flush small. + critical: bool = True def __init__( self, @@ -168,14 +214,26 @@ def __init__( self._tool_calls_count: int = 0 self._stopped_by: str = "" - # JSON envelope state. Messages are append-only on disk rather than - # retained in a per-agent list for the whole run. The public - # ``.json`` snapshot is streamed from this spool. - self._message_spool_handle: IO[str] | None = None - self._message_count: int = 0 - self._emitted_prior: bool = False - self._tool_call_ids: dict[int, list[str]] = {} - self._tool_results_seen: dict[int, int] = {} + # Borrow the loop's live history instead of reconstructing it from + # execution events. The spool is refreshed only when a snapshot is due. + # No second in-memory copy of the message bodies is retained. + self._history: list[Message] | None = None + # True when no loop history was offered (a host driving the hooks + # directly), so the observer assembles the conversation itself. + self._owns_history: bool = False + # Diagnostic fields are separate from history, kept as [message, + # content fingerprint, details] in response order. Identity binds them + # first; the fingerprint re-binds them to the shallow copies that + # compaction and thinking resets substitute. Entries whose message left + # history are released at the next snapshot. + self._response_metadata: list[list[Any]] = [] + # One (message, signature, end offset) per history entry already in + # the spool. A snapshot keeps the longest unchanged prefix and + # re-renders only from the first new, replaced or edited message. + self._spool_entries: list[tuple[Any, tuple[Any, ...], int]] = [] + # Set by on_loop_end. A cancellation arriving after it (while a later + # observer still runs) must not overwrite the terminal snapshot. + self._ended: bool = False # Flush coalescing cursors (see ``_flush_json``). ``0.0`` rather than # ``time.monotonic()`` so the FIRST flush is always due: the envelope # then appears on disk immediately instead of only after the first @@ -275,26 +333,143 @@ def _close_jsonl(self) -> None: # ── JSON envelope writer ──────────────────────────────────────────── - def _append_message(self, message: dict) -> None: - """Append one OpenAI message to the private on-disk spool.""" - if self._message_spool_handle is None: - self._message_spool_handle = open( # noqa: SIM115 - self._message_spool_path(), "w", encoding="utf-8", - ) - for chunk in _STREAM_ENCODER.iterencode(message): - self._message_spool_handle.write(chunk) - self._message_spool_handle.write("\n") - self._message_spool_handle.flush() - self._message_count += 1 - def _close_message_spool(self, *, cleanup: bool = False) -> None: - if self._message_spool_handle is not None: - self._message_spool_handle.close() - self._message_spool_handle = None + self._spool_entries = [] if cleanup: with contextlib.suppress(FileNotFoundError): self._message_spool_path().unlink() + @staticmethod + def _signature(message: Any, details: Any) -> tuple[Any, ...]: + """Shallow structure of *message*, compared by reference identity. + + History edits replace values (``message["content"] = rebuilt``) or + whole messages, so identity of each value and content block detects + them without re-serialising. Holding the references keeps their ids + from being reused before the next comparison. + """ + if not isinstance(message, dict): + return (message, details) + content = message.get("content") + blocks = tuple(content) if isinstance(content, list) else () + return (tuple(message.items()), blocks, details) + + @staticmethod + def _same(a: Any, b: Any) -> bool: + if isinstance(a, tuple) and isinstance(b, tuple): + return len(a) == len(b) and all( + TrajectoryFileObserver._same(x, y) for x, y in zip(a, b, strict=True) + ) + return a is b + + @staticmethod + def _fingerprint(message: Message) -> int: + """Content identity that survives reasoning and thinking-block strips.""" + content = message.get("content") + if isinstance(content, list): + content = [ + block for block in content + if not isinstance(block, dict) + or block.get("type") not in _THINKING_BLOCK_TYPES + ] + ids = tuple( + tc.get("id") for tc in (message.get("tool_calls") or []) + if isinstance(tc, dict) + ) + return hash((ids, _STREAM_ENCODER.encode(content))) + + def _resolve_annotations(self, history: list[Message]) -> dict[int, dict[str, Any]]: + """Bind diagnostics to the history messages that carry them now.""" + present = {id(message): message for message in history} + bound: dict[int, list[Any]] = {} + unmatched: list[list[Any]] = [] + for entry in self._response_metadata: + ref = entry[0] + if present.get(id(ref)) is ref and id(ref) not in bound: + bound[id(ref)] = entry + else: + unmatched.append(entry) + if unmatched: + pending: dict[int, deque[list[Any]]] = {} + for entry in unmatched: + pending.setdefault(entry[1], deque()).append(entry) + for message in history: + if not pending: + break + if ( + not isinstance(message, dict) + or message.get("role") != "assistant" + or id(message) in bound + ): + continue + fingerprint = self._fingerprint(message) + queue = pending.get(fingerprint) + if queue: + entry = queue.popleft() + entry[0] = message + bound[id(message)] = entry + if not queue: + del pending[fingerprint] + self._response_metadata = [ + entry for entry in self._response_metadata + if bound.get(id(entry[0])) is entry + ] + return {key: entry[2] for key, entry in bound.items()} + + def _render(self, message: Any, details: dict[str, Any] | None) -> bytes: + rendered = self._message_to_dict(message) + if rendered is None: + return b"" + if details: + blocks = details.get("thinking_blocks") + content = message.get("content") + if ( + blocks and isinstance(content, list) + and all(block in content for block in blocks) + ): + # Already preserved verbatim as native content blocks. + details = {k: v for k, v in details.items() if k != "thinking_blocks"} + rendered.update(details) + return (_STREAM_ENCODER.encode(rendered) + "\n").encode("utf-8") + + def _sync_history(self) -> None: + """Bring the spool up to the real history, including in-place edits. + + A length/delta cursor is insufficient: compaction, rollback, thinking + resets and image eviction can replace or edit already-recorded messages. + The unchanged prefix stays on disk; from the first difference onwards + the spool is truncated and re-rendered. Each message is rendered whole + before it is written, so a failure leaves a consistent prefix. + """ + history = self._history + if history is None: + return + details_for = self._resolve_annotations(history) + entries = self._spool_entries + path = self._message_spool_path() + keep = 0 + if path.exists(): + limit = min(len(entries), len(history)) + while keep < limit: + message = history[keep] + ref, signature, _ = entries[keep] + current = self._signature(message, details_for.get(id(message))) + if ref is not message or not self._same(signature, current): + break + keep += 1 + offset = entries[keep - 1][2] if keep else 0 + del entries[keep:] + # Open even for an empty history so a rollback to [] clears the spool. + with path.open("r+b" if keep else "wb") as fh: + fh.seek(offset) + fh.truncate() + for message in history[keep:]: + details = details_for.get(id(message)) + line = self._render(message, details) + fh.write(line) + offset += len(line) + entries.append((message, self._signature(message, details), offset)) + def _write_envelope(self, path: Path) -> None: """Stream one compatible JSON envelope without materialising it.""" fields: list[tuple[str, Any]] = [ @@ -334,11 +509,16 @@ def _write_envelope(self, path: Path) -> None: fh.write(line.rstrip("\n")) fh.write("]}") + def _count_change(self) -> None: + """Count a history change toward the next coalesced snapshot.""" + if "json" in self._formats: + self._pending_flushes += 1 + def _flush_json(self, *, force: bool = False) -> None: """Rewrite the JSON envelope, coalescing bursts unless *force*. The envelope is a single document, so every flush still copies all - message chunks from the append-only spool. Coalescing bounds that + message chunks from the refreshed history spool. Coalescing bounds that O(n²) I/O to one copy per ``_COALESCE_N`` events or ``_COALESCE_MS`` milliseconds. The copy is streamed, so neither the message history nor its serialised twin is materialised in memory. @@ -373,27 +553,12 @@ def _flush_json(self, *, force: bool = False) -> None: return self._pending_flushes = 0 self._last_flush_at = time.monotonic() + self._sync_history() path = self._path("json") tmp = path.with_suffix(path.suffix + ".tmp") self._write_envelope(tmp) os.replace(tmp, path) - @staticmethod - def _synth_id(turn: int, idx: int) -> str: - return f"call_{turn}_{idx}" - - @staticmethod - def _stringify(value: Any) -> str: - if isinstance(value, str): - return value - return str(value) if value is not None else "" - - @staticmethod - def _format_args(args: Any) -> str: - if isinstance(args, (dict, list)): - return json.dumps(args, ensure_ascii=False) - return TrajectoryFileObserver._stringify(args) - def _message_to_dict(self, m: Any) -> dict | None: """Pass through native OpenAI-shaped dict messages. @@ -406,7 +571,19 @@ def _message_to_dict(self, m: Any) -> dict | None: # ``redacted_for_trace`` swaps an inline image's base64 for its # size. Verbatim, one 1080p screenshot writes ~137 KB here for # every turn it stays in history. - return dict(redacted_for_trace(m)) + message = dict(redacted_for_trace(m)) + if message.get("role") == "tool": + content = message.get("content") + if isinstance(content, str): + message["content"] = _clip(content, _BODY_MAX_CHARS) + elif isinstance(content, list): + message["content"] = [ + {**block, "text": _clip(block["text"], _BODY_MAX_CHARS)} + if isinstance(block, dict) and isinstance(block.get("text"), str) + else block + for block in content + ] + return message return None # ── Lifecycle hooks ───────────────────────────────────────────────── @@ -419,13 +596,11 @@ def _message_to_dict(self, m: Any) -> dict | None: # ``recover_result`` needs to find this file, and the file is not # sandbox-visible, so the path is published rather than mounted. # - # Published into ``ExecutionScope.metadata`` and NOT into a contextvar. This - # hook is dispatched by ``notify_observers`` as ``asyncio.create_task`` for - # non-critical observers, and this observer is non-critical — a contextvar - # ``.set()`` inside that task mutates only the task's own copy of the context - # and is invisible to the loop. The scope OBJECT, by contrast, is shared by - # reference through the inherited context, so mutating its dict here is seen - # by the loop and by every nested tool-call task. Same channel + # Published into ``ExecutionScope.metadata`` and NOT into a contextvar. + # Non-critical hooks may run in ``asyncio.create_task``; a contextvar + # ``.set()`` inside such a task mutates only that task's own context. Scope + # metadata works for both awaited and background observers, and is seen + # by the loop and every nested tool-call task. Same channel # ``agent_loop.py`` already uses for ``current_turn``. SCOPE_KEY = "trajectory_jsonl" @@ -463,6 +638,7 @@ async def on_loop_start(self, config: LoopConfig) -> None: self._task_id = config.task_id self._role_id = config.role_id self._max_turns = config.max_turns + self._ended = False # After ``_task_id`` is set: it is the stem fallback when no explicit # ``filename`` was passed. self._publish_jsonl_path() @@ -484,6 +660,65 @@ async def on_loop_start(self, config: LoopConfig) -> None: self._write_jsonl(start_record) self._flush_json() + async def on_before_llm(self, ctx: TurnContext) -> Intervention | None: + # Unlike on_llm_input's request copy, this list stays live across + # injections, compaction, early exits and cancellation. The loop + # rewrites it in place, so binding once per run is enough. + self._bind_history(ctx.messages) + return None + + async def on_llm_input(self, ctx: TurnContext) -> None: + # Injections from on_before_llm have now been applied. Transient request + # addenda belong to the request, not to durable conversation history. + self._flush_json() + + async def on_turn_end(self, ctx: TurnContext) -> Intervention | None: + self._flush_json() + return None + + def _bind_history(self, messages: list[Message] | None) -> None: + if "json" not in self._formats or self._owns_history or messages is None: + return + if self._history is None: + self._history = messages + + @staticmethod + def _assistant_from_ctx(ctx: TurnContext) -> Message: + """Build the assistant entry for hosts that supply no loop history.""" + message: dict[str, Any] = {"role": "assistant", "content": ctx.ai_text or ""} + if ctx.thinking: + message["reasoning_content"] = ctx.thinking + if ctx.tool_calls: + message["tool_calls"] = [ + { + "id": tc.get("id") or f"call_{ctx.turn}_{idx}", + "type": "function", + "function": { + "name": tc.get("name", ""), + "arguments": ( + _STREAM_ENCODER.encode(tc.get("args", {})) + if isinstance(tc.get("args", {}), (dict, list)) + else str(tc.get("args") or "") + ), + }, + } + for idx, tc in enumerate(ctx.tool_calls) + ] + return message # type: ignore[return-value] + + def _owned_tool_call_id(self, result: ToolResult) -> str: + cid = getattr(result, "tool_call_id", "") or "" + if cid or self._history is None: + return cid + answered = {m.get("tool_call_id") for m in self._history if m.get("role") == "tool"} + for message in reversed(self._history): + if message.get("role") == "assistant": + for tc in message.get("tool_calls") or []: + if tc["id"] not in answered: + return tc["id"] + break + return "" + async def on_llm_response( self, ctx: TurnContext, ) -> Intervention | None: @@ -520,57 +755,32 @@ async def on_llm_response( record["stop_details"] = ctx.stop_details self._write_jsonl(record) - # JSON envelope: emit prior context once, then this turn. if "json" in self._formats: - if not self._emitted_prior: - # ctx.messages includes the assistant we just got; skip it - # and emit ourselves below with reasoning_content split out. - for m in list(ctx.messages or [])[:-1]: - d = self._message_to_dict(m) - if d: - self._append_message(d) - self._emitted_prior = True - - msg: dict = { - "role": "assistant", - "content": self._stringify(ctx.ai_text), - } - if ctx.thinking: - msg["reasoning_content"] = self._stringify(ctx.thinking) - if ctx.usage: - # Per-turn usage (incl. reasoning_tokens) so the JSON envelope - # carries the same usage as the .jsonl stream. - msg["usage"] = ctx.usage - if ctx.thinking_blocks: - # Verbatim thinking / reasoning blocks (signatures / - # encrypted_content) so the JSON envelope stays replay-able. - msg["thinking_blocks"] = ctx.thinking_blocks - if ctx.finish_reason and ctx.finish_reason not in _ORDINARY_FINISH_REASONS: - msg["finish_reason"] = ctx.finish_reason - if ctx.stop_details: - msg["stop_details"] = ctx.stop_details - if ctx.tool_calls: - ids: list[str] = [] - tcs: list[dict] = [] - for idx, tc in enumerate(ctx.tool_calls): - cid = tc.get("id") or self._synth_id(ctx.turn, idx) - ids.append(cid) - tcs.append({ - "id": cid, - "type": "function", - "function": { - "name": tc.get("name", ""), - "arguments": self._format_args( - tc.get("args", {}), - ), - }, - }) - msg["tool_calls"] = tcs - self._tool_call_ids[ctx.turn] = ids - self._tool_results_seen[ctx.turn] = 0 - self._tool_calls_count += len(tcs) - - self._append_message(msg) + messages = ctx.messages or [] + self._bind_history(messages if messages else None) + # The last assistant may already be followed by skipped/blocked + # replies. Never infer its position from messages[-1]. Awaited + # dispatch ensures this still refers to the response being observed. + assistant = next( + (m for m in reversed(messages) + if isinstance(m, dict) and m.get("role") == "assistant"), + None, + ) + if assistant is None and (self._history is None or self._owns_history): + # No loop history at all: keep the pre-history-snapshot + # contract and record the turn from the context itself. + self._owns_history = True + if self._history is None: + self._history = [] + assistant = self._assistant_from_ctx(ctx) + self._history.append(assistant) + if assistant is not None: + details = {key: record[key] for key in _DIAGNOSTIC_KEYS if key in record} + if details: + self._response_metadata.append( + [assistant, self._fingerprint(assistant), details], + ) + self._tool_calls_count += len(ctx.tool_calls or []) self._turns_used = ctx.turn self._flush_json() return None @@ -578,13 +788,9 @@ async def on_llm_response( async def on_tool_result( self, ctx: TurnContext, result: ToolResult, ) -> ToolResult | None: - # ``tool_call_id`` is what makes a JSONL entry addressable: a turn can - # hold several results from the same tool (parallel_tool_calls is on), - # so ``(turn, name)`` does not identify one. Recorded straight off the - # result rather than through the JSON branch's synthesis fallback below - # — that fallback advances ``_tool_results_seen``, so sharing it would - # double-count, and a synthesised id matches nothing outside the - # snapshot anyway. Empty here means the runtime itself had no id. + # Results describe what the tool returned, before post-processing and + # image capability gating. Only the real history can say what the model + # received. Empty ids here mean the runtime itself supplied no id. jsonl_record: dict[str, Any] = { "t": "result", "turn": ctx.turn, @@ -601,26 +807,23 @@ async def on_tool_result( self._write_jsonl(jsonl_record) if "json" in self._formats: - cid = getattr(result, "tool_call_id", "") or "" - if not cid: - ids = self._tool_call_ids.get(ctx.turn, []) - seen = self._tool_results_seen.get(ctx.turn, 0) - cid = ( - ids[seen] if seen < len(ids) - else self._synth_id(ctx.turn, seen) + if self._owns_history and self._history is not None: + body = _clip( + result.result if isinstance(result.result, str) + else ("" if result.result is None else str(result.result)), + _BODY_MAX_CHARS, ) - self._tool_results_seen[ctx.turn] = seen + 1 - body = _clip(self._stringify(result.result), _BODY_MAX_CHARS) - rendered_body = f"[error] {body}" if result.is_error else body - self._append_message({ - "role": "tool", - "tool_call_id": cid, - "content": redacted_tool_result_content( - rendered_body, - result.images, - ), - }) - self._flush_json() + self._history.append({ # type: ignore[arg-type] + "role": "tool", + "tool_call_id": self._owned_tool_call_id(result), + "content": redacted_tool_result_content( + f"[error] {body}" if result.is_error else body, + result.images, + ), + }) + # The loop appends the reply only after this hook, so a snapshot + # here would show an unanswered call. Count it toward the next one. + self._count_change() return None async def on_compaction(self, event: CompactionEvent) -> None: @@ -647,6 +850,9 @@ async def on_compaction(self, event: CompactionEvent) -> None: }) async def on_loop_end(self, result: AgentLoopResult) -> None: + if "json" in self._formats and not self._owns_history: + self._history = result.messages + self._ended = True self._turns_used = max(self._turns_used, result.turns_used) self._tool_calls_count = result.tool_calls_count self._stopped_by = ( @@ -673,9 +879,22 @@ async def on_loop_end(self, result: AgentLoopResult) -> None: # Always release the FD. Keep the forensic spool when the terminal # snapshot fails; remove it only after a successful atomic replace. self._close_message_spool(cleanup=snapshot_written) + self._history = None + self._owns_history = False + self._response_metadata = [] async def on_loop_cancelled(self) -> None: - """Close live handles while preserving crash-forensic sidecars.""" + """Flush current history and preserve crash-forensic sidecars.""" self._close_jsonl() self._withdraw_jsonl_path() - self._close_message_spool() + if self._ended: + # on_loop_end already wrote the terminal snapshot and released the + # history; re-flushing would replace it with an empty envelope. + return + try: + self._flush_json(force=True) + finally: + self._close_message_spool() + self._history = None + self._owns_history = False + self._response_metadata = [] diff --git a/changes/83.fix.md b/changes/83.fix.md new file mode 100644 index 0000000..241abf4 --- /dev/null +++ b/changes/83.fix.md @@ -0,0 +1 @@ +Fix trajectory JSON snapshots to follow the loop's actual message history, including capped and blocked native tool calls, synthetic replies, post-processed tool bodies, observer injections, rollbacks and compaction. Await trajectory hooks so immediate invalid-call results cannot precede their assistant event. Preserve normalizer-controlled reasoning and native content blocks, diagnostic fields, image redaction, configured body limits and coalesced atomic writes; flush the current history on completion or cancellation. JSONL remains the event log with original tool outputs, including events later removed from conversation history. Snapshots re-render only new or changed messages, keep diagnostics through compaction copies, tolerate non-JSON values, and are not overwritten by a cancellation after completion. diff --git a/docs/agent-loop-boundary.md b/docs/agent-loop-boundary.md index b49c6ef..7ebe1b4 100644 --- a/docs/agent-loop-boundary.md +++ b/docs/agent-loop-boundary.md @@ -40,6 +40,37 @@ and falls back to core behavior rather than failing the batch, and a fan-in interrupt waiter that returns `False` or raises leaves the tool running instead of discarding its collected work. +## Trajectory snapshots and events + +`TrajectoryFileObserver` writes two complementary formats. The JSON envelope's +`messages` follows the loop's current history, including engine-generated tool +replies, post-processed result bodies, injections, rollbacks and compaction. +It preserves the normalizer's reasoning policy and native content blocks rather +than reconstructing assistant messages from the executable subset of calls. +Diagnostic fields (`usage`, `thinking_blocks`, `finish_reason`, `stop_details`) +remain annotations on the corresponding assistant message, including after +compaction or a thinking reset substitutes a copy of it; `thinking_blocks` is +omitted while the same blocks are still native assistant content. Values that +are not JSON (datetimes, bytes, SDK objects) are written in a lossy string form +rather than failing the snapshot. Images are redacted +and tool text obeys `SWARM_TRAJECTORY_BODY_MAX_CHARS` (default 16384; 0 disables +the cap), so the snapshot is a trace representation rather than a byte-exact +provider payload. Transient per-request addenda and provider middleware edits +are outside this durable-history snapshot. + +JSON snapshots are coalesced during a run (LLM input and response, each tool +result and turn end count toward `SWARM_TRAJECTORY_COALESCE_N`) and atomically +replaced; each snapshot re-renders only messages added or changed since the +previous one. Completion and cancellation force a snapshot of the latest +history, and a cancellation that arrives after completion leaves the terminal +snapshot untouched. Hosts that drive the hooks without loop history still get +the assistant and tool entries built from the hook arguments. JSONL is the event +log: it retains original tool outputs before post-processing, and events whose +messages are later pruned from history. Trajectory hooks are awaited so an +immediate synthetic result cannot overtake its assistant event. Consumers that +need the execution events should read JSONL; consumers that need the current +conversation should read JSON and handle native block-list content. + ## Repeated-invocation metadata has no core-side consumer, by construction `ToolResult.repeat_count`, `ToolResult.repeat_recovery_id`, `ToolResult.result_id` diff --git a/tests/test_tool_content_images.py b/tests/test_tool_content_images.py index 1e7cc35..4a8aebb 100644 --- a/tests/test_tool_content_images.py +++ b/tests/test_tool_content_images.py @@ -535,12 +535,14 @@ def test_the_trajectory_observer_does_not_write_base64() -> None: @pytest.mark.asyncio -async def test_live_trajectory_records_returned_images_without_base64(tmp_path) -> None: +@pytest.mark.parametrize("visible", [False, True]) +async def test_live_trajectory_records_image_delivery_without_base64(tmp_path, visible) -> None: import json - from types import SimpleNamespace from agent_core.components.observers.trajectory import TrajectoryFileObserver - from agent_core.loop_types import LoopConfig, ToolResult + from agent_core.llm import LLMResponse + from agent_core.loop_types import LoopConfig, LoopPolicy + from agent_core.runtime.loop.agent_loop import run_agent_loop attachment = image_attachment( _png(64, 32), @@ -552,25 +554,41 @@ async def test_live_trajectory_records_returned_images_without_base64(tmp_path) filename="live", formats=["json", "jsonl"], ) - await observer.on_loop_start(LoopConfig(task_id="probe")) - await observer.on_tool_result( - SimpleNamespace(turn=1), - ToolResult( - name="view_image", - args={"path": "/tmp/live.png"}, - result="image follows", - duration_ms=1, - tool_call_id="call_1", - is_error=False, - images=[attachment], - ), + class ViewImage: + name = "view_image" + + async def ainvoke(self, args): + return tool_content("image follows", images=[attachment]) + + class Client: + called = False + + async def chat(self, messages, **kwargs): + if not self.called: + self.called = True + return LLMResponse(tool_calls=[{ + "id": "call_1", "type": "function", "function": { + "name": "view_image", "arguments": '{"path":"/tmp/live.png"}', + }, + }]) + return LLMResponse(content="done") + + await run_agent_loop( + system_prompt="system", user_message="show image", llm=Client(), + tools=[ViewImage()], observers=[observer], + model_profile=_vision_profile() if visible else _blind_profile(), + config=LoopConfig(task_id="probe", max_turns=2, + loop_policy=LoopPolicy(no_tool_behavior="stop")), ) - observer._flush_json(force=True) envelope = json.loads((tmp_path / "live.json").read_text(encoding="utf-8")) - content = envelope["messages"][-1]["content"] - assert content[1]["type"] == "image_url" - assert "elided from trace" in content[1]["image_url"]["url"] + content = next(m["content"] for m in envelope["messages"] if m["role"] == "tool") + if visible: + assert content[1]["type"] == "image_url" + assert "elided from trace" in content[1]["image_url"]["url"] + else: + assert isinstance(content, str) + assert "have NOT" in content jsonl = (tmp_path / "live.jsonl").read_text(encoding="utf-8") assert '"images"' in jsonl diff --git a/tests/test_trajectory_history.py b/tests/test_trajectory_history.py new file mode 100644 index 0000000..f0f90c7 --- /dev/null +++ b/tests/test_trajectory_history.py @@ -0,0 +1,469 @@ +"""The JSON trajectory follows the loop's history, including engine replies.""" + +from __future__ import annotations + +import asyncio +import copy +import json + +import pytest + +from agent_core.components.observers.last_turn_forcer import LastTurnForcer +from agent_core.components.observers.trajectory import TrajectoryFileObserver +from agent_core.llm import LLMResponse +from agent_core.loop_types import ( + AgentLoopResult, + BaseObserver, + Intervention, + LoopConfig, + LoopPolicy, + TurnContext, +) +from agent_core.runtime.loop.agent_loop import run_agent_loop +from agent_core.runtime.loop.model_profile import HistoryPolicy, ModelProfile + + +class EchoTool: + name = "echo" + + def to_openai_schema(self): + return {"type": "function", "function": { + "name": self.name, "parameters": {"type": "object"}, + }} + + async def ainvoke(self, args): + return args.get("value", "ok") + + +def _call(cid, arguments='{"value":"ok"}', name="echo"): + return {"id": cid, "type": "function", "function": { + "name": name, "arguments": arguments, + }} + + +def _messages(path): + return json.loads(path.read_text())["messages"] + + +def _without_diagnostics(messages): + return [{k: v for k, v in msg.items() + if k not in {"usage", "thinking_blocks", "finish_reason", "stop_details"}} + for msg in messages] + + +def _assert_balanced(messages): + pending = set() + for msg in messages: + if msg["role"] == "tool": + assert msg["tool_call_id"] in pending + pending.remove(msg["tool_call_id"]) + else: + assert not pending + pending.update(tc["id"] for tc in msg.get("tool_calls", [])) + assert not pending + + +async def _run(tmp_path, responses, *, observers=(), config=None, + formats=("json", "jsonl"), verify_snapshots=True, **kwargs): + class Client: + def __init__(self): + self.requests = [] + self.snapshots = [] + + async def chat(self, messages, **_kwargs): + self.requests.append(copy.deepcopy(messages)) + self.snapshots.append(_messages(tmp_path / "trace.json")) + return responses.pop(0) + + client = Client() + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=formats) + result = await run_agent_loop( + system_prompt="system", user_message="start", llm=client, tools=[EchoTool()], + observers=[trajectory, *observers], + config=config or LoopConfig(max_turns=5, max_llm_retries=1, + loop_policy=LoopPolicy(no_tool_behavior="stop")), + **kwargs, + ) + messages = _messages(tmp_path / "trace.json") + assert _without_diagnostics(messages) == result.messages + _assert_balanced(messages) + if verify_snapshots: + for snapshot, request in zip(client.snapshots, client.requests, strict=True): + assert _without_diagnostics(snapshot) == request + assert not (tmp_path / "trace.messages.spool").exists() + assert trajectory._history is None + assert not trajectory._response_metadata + return messages, client, result + + +@pytest.fixture(autouse=True) +def _uncoalesced(monkeypatch): + monkeypatch.setattr("agent_core.components.observers.trajectory._COALESCE_N", 0) + monkeypatch.setattr("agent_core.components.observers.trajectory._COALESCE_MS", 0) + + +@pytest.mark.parametrize("turn,call_count", [(1, 6), (1, 7), (2, 6)]) +@pytest.mark.parametrize("formats", [("json",), ("json", "jsonl")]) +async def test_over_cap_keeps_every_native_call_and_reply(tmp_path, turn, call_count, formats): + responses = ([LLMResponse(tool_calls=[_call("warmup")])] if turn == 2 else []) + responses += [ + LLMResponse(tool_calls=[_call(f"c{i}") for i in range(call_count)]), + LLMResponse(content="done"), + ] + messages, _, _ = await _run(tmp_path, responses, formats=formats, config=LoopConfig( + max_turns=4, max_llm_retries=1, max_tool_calls_per_turn=5, + loop_policy=LoopPolicy(no_tool_behavior="stop"), + )) + assistants = [m for m in messages if m["role"] == "assistant"] + assert len(assistants) == turn + 1 + assert len(assistants[turn - 1]["tool_calls"]) == call_count + skipped = [m for m in messages if m["role"] == "tool" and "[tool call skipped]" in m["content"]] + assert len(skipped) == call_count - 5 + assert (tmp_path / "trace.jsonl").exists() == ("jsonl" in formats) + + +@pytest.mark.parametrize("scenario", ["unknown", "final", "skip", "restart"]) +async def test_engine_generated_tool_replies_are_recorded(tmp_path, scenario): + class Control(BaseObserver): + critical = True + fired = False + + async def on_llm_response(self, ctx): + if ctx.turn == 2 and not self.fired: + self.fired = True + if scenario == "skip": + return Intervention(skip_tool_execution=True) + if scenario == "restart": + return Intervention(continue_to_next_turn=True) + return None + + responses = [ + LLMResponse(tool_calls=[_call("warmup")]), + LLMResponse(tool_calls=[_call("target", name="missing" if scenario == "unknown" else "echo"), + *([_call("known")] if scenario == "unknown" else [])]), + ] + if scenario != "final": + responses.append(LLMResponse(content="done")) + messages, _, _ = await _run(tmp_path, responses, observers=[Control(), LastTurnForcer()], config=LoopConfig( + max_turns=2 if scenario == "final" else 5, + max_llm_retries=1, + loop_policy=LoopPolicy(no_tool_behavior="stop"), + )) + target = next(m for m in messages if m.get("tool_call_id") == "target") + assert "[tool call" in target["content"] + + +async def test_invalid_arguments_follow_their_assistant_even_without_an_async_yield(tmp_path): + messages, _, _ = await _run(tmp_path, [ + LLMResponse(content="partial", tool_calls=[_call("broken", "{")], finish_reason="length"), + LLMResponse(content="done"), + ]) + assistant_index = next(i for i, m in enumerate(messages) if m.get("tool_calls")) + assert messages[assistant_index + 1]["tool_call_id"] == "broken" + assert "[invalid tool call]" in messages[assistant_index + 1]["content"] + assert messages[assistant_index]["finish_reason"] == "length" + events = [json.loads(line) for line in (tmp_path / "trace.jsonl").read_text().splitlines()] + assert [e["t"] for e in events] == ["start", "llm", "result", "llm", "end"] + + +@pytest.mark.parametrize("rollback", [False, True]) +async def test_injections_rollbacks_and_compaction_follow_rewritten_history(tmp_path, rollback): + class Rewrite(BaseObserver): + critical = True + fired = False + + async def on_turn_end(self, ctx): + if not self.fired: + self.fired = True + return Intervention(inject_messages=["continue"], pop_last_message=rollback, + continue_to_next_turn=rollback) + return None + + class Compact: + def compact(self, messages, keep_recent): + return [messages[0], {"role": "user", "content": "compacted"}] + + messages, _, _ = await _run(tmp_path, [ + LLMResponse(tool_calls=[_call("discarded")]), LLMResponse(content="done"), + ], observers=[Rewrite()], config=LoopConfig( + max_turns=4, max_llm_retries=1, compact_after_turns=0, + compactor=Compact(), loop_policy=LoopPolicy(no_tool_behavior="stop"), + )) + assert "discarded" not in json.dumps(messages) + assert any(m["role"] == "user" for m in messages) + + +async def test_tool_snapshot_uses_postprocessed_body_and_jsonl_keeps_original(tmp_path): + class Processor: + def process(self, result): + return "model preview" + + messages, _, _ = await _run(tmp_path, [ + LLMResponse(tool_calls=[_call("c", '{"value":"original result"}')]), + LLMResponse(content="done"), + ], config=LoopConfig(max_turns=3, tool_result_post_processor=Processor(), + loop_policy=LoopPolicy(no_tool_behavior="stop"))) + assert next(m for m in messages if m["role"] == "tool")["content"] == "model preview" + events = [json.loads(line) for line in (tmp_path / "trace.jsonl").read_text().splitlines()] + assert next(e for e in events if e["t"] == "result")["result"] == "original result" + + +@pytest.mark.parametrize("thinking_format", ["tag", "reasoning_content", "content_block"]) +@pytest.mark.parametrize("thinking_in_history", [False, True]) +async def test_reasoning_uses_the_normalizers_history_policy(tmp_path, thinking_format, thinking_in_history): + content = ([{"type": "thinking", "thinking": "reasoning", "signature": "signed"}, + {"type": "text", "text": "visible"}] + if thinking_format == "content_block" else "visible") + messages, _, _ = await _run(tmp_path, [ + LLMResponse(content=content, reasoning_content="reasoning", tool_calls=[_call("c")]), + LLMResponse(content="done"), + ], model_profile=ModelProfile(model_id="test", provider="test", thinking_format=thinking_format), + history_policy=HistoryPolicy(thinking_in_history=thinking_in_history)) + first = next(m for m in messages if m["role"] == "assistant") + if thinking_format == "content_block": + assert first["content"] == content + elif thinking_format == "reasoning_content" and thinking_in_history: + assert first["reasoning_content"] == "reasoning" + else: + assert "reasoning_content" not in first + + +@pytest.mark.parametrize("early_stop", [False, True]) +async def test_empty_run_and_before_llm_stop_keep_initial_history(tmp_path, early_stop): + class Stop(BaseObserver): + critical = True + + async def on_before_llm(self, ctx): + return Intervention(stop_reason="stopped", inject_messages=["notice"]) + + messages, client, _ = await _run(tmp_path, [], observers=[Stop()], config=LoopConfig( + max_turns=3 if early_stop else 0, + )) + assert not client.requests + assert [m["content"] for m in messages] == ["system", "start", *(["notice"] if early_stop else [])] + + +async def test_coalesced_terminal_snapshot_is_complete(tmp_path, monkeypatch): + monkeypatch.setattr("agent_core.components.observers.trajectory._COALESCE_N", 1000) + messages, _, _ = await _run(tmp_path, [ + LLMResponse(tool_calls=[_call(f"c{i}") for i in range(7)]), + LLMResponse(content="done"), + ], verify_snapshots=False) + assert len([m for m in messages if m["role"] == "tool"]) == 7 + + +async def test_cancellation_flushes_current_history_and_preserves_spool(tmp_path): + entered = asyncio.Event() + + class Client: + async def chat(self, messages, **kwargs): + entered.set() + await asyncio.Event().wait() + + trajectory = TrajectoryFileObserver(tmp_path, filename="trace") + task = asyncio.create_task(run_agent_loop( + system_prompt="system", user_message="start", llm=Client(), tools=[], + observers=[trajectory], config=LoopConfig(max_turns=2), + )) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert [m["content"] for m in _messages(tmp_path / "trace.json")] == ["system", "start"] + assert (tmp_path / "trace.messages.spool").exists() + + +async def test_failed_terminal_replace_keeps_readable_snapshot_and_forensic_spool(tmp_path, monkeypatch): + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=["json"]) + await trajectory.on_loop_start(LoopConfig()) + previous = (tmp_path / "trace.json").read_bytes() + + def fail_replace(source, destination): + raise OSError("disk failure") + + monkeypatch.setattr("agent_core.components.observers.trajectory.os.replace", fail_replace) + with pytest.raises(OSError, match="disk failure"): + await trajectory.on_loop_end(AgentLoopResult(messages=[{"role": "user", "content": "latest"}])) + assert (tmp_path / "trace.json").read_bytes() == previous + assert json.loads((tmp_path / "trace.messages.spool").read_text())["content"] == "latest" + assert trajectory._spool_entries == [] + assert trajectory._history is None + + +@pytest.mark.parametrize("formats", [("jsonl",), ()]) +async def test_disabled_json_creates_no_snapshot_or_spool(tmp_path, formats): + class Client: + async def chat(self, messages, **kwargs): + return LLMResponse(content="done") + + await run_agent_loop( + system_prompt="system", user_message="start", llm=Client(), tools=[], + observers=[TrajectoryFileObserver(tmp_path, filename="trace", formats=formats)], + config=LoopConfig(max_turns=1, loop_policy=LoopPolicy(no_tool_behavior="stop")), + ) + assert not (tmp_path / "trace.json").exists() + assert not (tmp_path / "trace.messages.spool").exists() + assert (tmp_path / "trace.jsonl").exists() == ("jsonl" in formats) + + +@pytest.mark.parametrize("limit", [0, 4]) +def test_history_body_limits_preserve_blocks_and_do_not_mutate_history(tmp_path, monkeypatch, limit): + monkeypatch.setattr("agent_core.components.observers.trajectory._BODY_MAX_CHARS", limit) + message = {"role": "tool", "tool_call_id": "c", "content": [ + {"type": "text", "text": "long result"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,payload"}}, + ]} + original = copy.deepcopy(message) + trajectory = TrajectoryFileObserver(tmp_path) + rendered = trajectory._message_to_dict(message) + text = rendered["content"][0]["text"] + if limit == 0: + assert text == "long result" + else: + assert text.startswith("long\n... [truncated") + assert "elided from trace" in rendered["content"][1]["image_url"]["url"] + assert message == original + + +def _turn(messages, **kwargs): + fields = {"turn": 1, "max_turns": 5, "task_id": "t", "role_id": "r", "ai_text": "", + "thinking": "", "tool_calls": [], "messages": messages, "usage": None, + "metadata": {}} + fields.update(kwargs) + return TurnContext(**fields) + + +async def test_cancellation_after_loop_end_keeps_terminal_snapshot(tmp_path): + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=["json"]) + await trajectory.on_loop_start(LoopConfig()) + history = [{"role": "user", "content": "start"}, {"role": "assistant", "content": "done"}] + await trajectory.on_loop_end(AgentLoopResult(messages=history)) + terminal = (tmp_path / "trace.json").read_bytes() + # A later critical observer's on_loop_end is cancelled; the loop then + # notifies on_loop_cancelled on every observer, this one included. + await trajectory.on_loop_cancelled() + assert (tmp_path / "trace.json").read_bytes() == terminal + assert [m["content"] for m in _messages(tmp_path / "trace.json")] == ["start", "done"] + + +async def test_diagnostics_survive_compaction_that_copies_assistants(tmp_path): + class CopyCompact: + def compact(self, messages, keep_recent): + # Like strip_stale_reasoning: same content, new dict objects. + return [{k: v for k, v in m.items() if k != "reasoning_content"} for m in messages] + + messages, _, _ = await _run(tmp_path, [ + LLMResponse(tool_calls=[_call("c")], finish_reason="tool_use"), + LLMResponse(content="done"), + ], config=LoopConfig(max_turns=4, max_llm_retries=1, compact_after_turns=0, + compactor=CopyCompact(), + loop_policy=LoopPolicy(no_tool_behavior="stop"))) + first = next(m for m in messages if m["role"] == "assistant") + assert first["finish_reason"] == "tool_use" + + +async def test_diagnostics_survive_a_signed_thinking_reset(tmp_path): + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=["json"]) + blocks = [{"type": "thinking", "thinking": "r", "signature": "s"}] + assistant = {"role": "assistant", "content": [*blocks, {"type": "text", "text": "a"}]} + history = [{"role": "user", "content": "start"}, assistant] + await trajectory.on_llm_response(_turn(history, ai_text="a", thinking_blocks=blocks, + usage={"input_tokens": 3})) + # _reset_signed_thinking_history substitutes a copy without the blocks. + history[1] = {**assistant, "content": [{"type": "text", "text": "a"}]} + await trajectory.on_loop_end(AgentLoopResult(messages=history)) + last = _messages(tmp_path / "trace.json")[-1] + assert last["usage"] == {"input_tokens": 3} + # The blocks are gone from content, so the annotation must carry them. + assert last["thinking_blocks"] == blocks + + +async def test_native_thinking_blocks_are_not_written_twice(tmp_path): + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=["json"]) + blocks = [{"type": "thinking", "thinking": "r", "signature": "s"}] + history = [{"role": "assistant", "content": [*blocks, {"type": "text", "text": "a"}]}] + await trajectory.on_llm_response(_turn(history, thinking_blocks=blocks, + usage={"input_tokens": 3})) + message = _messages(tmp_path / "trace.json")[0] + assert "thinking_blocks" not in message + assert message["usage"] == {"input_tokens": 3} + + +async def test_non_json_history_values_do_not_freeze_snapshots(tmp_path): + import datetime + + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=["json", "jsonl"]) + await trajectory.on_loop_start(LoopConfig()) + stamp = datetime.datetime(2026, 1, 2, tzinfo=datetime.UTC) + history = [ + {"role": "assistant", "content": "", "tool_calls": [_call("c")]}, + {"role": "tool", "tool_call_id": "c", "content": "ok", "seen_at": stamp, + "raw": b"abc"}, + ] + await trajectory.on_loop_end(AgentLoopResult(messages=history)) + tool = _messages(tmp_path / "trace.json")[-1] + assert tool["seen_at"] == stamp.isoformat() + assert tool["raw"] == "<3 bytes>" + + +async def test_hooks_without_loop_history_still_record_the_turn(tmp_path): + from agent_core.loop_types import ToolResult + + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=["json"]) + await trajectory.on_loop_start(LoopConfig()) + await trajectory.on_llm_response(_turn( + None, ai_text="calling", tool_calls=[{"name": "echo", "args": {"v": 1}}], + finish_reason="tool_use", + )) + await trajectory.on_tool_result(_turn(None), ToolResult( + name="echo", args={"v": 1}, result="ok", duration_ms=1, tool_call_id="", is_error=False, + )) + await trajectory.on_llm_response(_turn(None, turn=2, ai_text="done", tool_calls=None)) + await trajectory.on_loop_end(AgentLoopResult(messages=[], turns_used=2)) + messages = _messages(tmp_path / "trace.json") + assert [m["role"] for m in messages] == ["assistant", "tool", "assistant"] + assert messages[0]["finish_reason"] == "tool_use" + assert messages[0]["tool_calls"][0]["function"]["arguments"] == '{"v":1}' + assert messages[1]["tool_call_id"] == messages[0]["tool_calls"][0]["id"] + assert messages[2]["content"] == "done" + + +async def test_tool_results_count_toward_coalesced_snapshots(tmp_path, monkeypatch): + monkeypatch.setattr("agent_core.components.observers.trajectory._COALESCE_N", 5) + _, client, _ = await _run(tmp_path, [ + LLMResponse(tool_calls=[_call(f"c{i}") for i in range(3)]), + LLMResponse(content="done"), + ], verify_snapshots=False) + # Turn 1 is input + response + 3 results + turn end; without the results + # it would only become due at the second input, after the snapshot. + assert [m.get("tool_call_id") for m in client.snapshots[1] if m["role"] == "tool"] == [ + "c0", "c1", "c2", + ] + + +async def test_snapshots_render_only_new_or_changed_messages(tmp_path, monkeypatch): + trajectory = TrajectoryFileObserver(tmp_path, filename="trace", formats=["json"]) + rendered = [] + original = trajectory._message_to_dict + + def counting(message): + rendered.append(message.get("content")) + return original(message) + + monkeypatch.setattr(trajectory, "_message_to_dict", counting) + history = [{"role": "user", "content": f"m{i}"} for i in range(3)] + trajectory._history = history + trajectory._flush_json(force=True) + history.append({"role": "user", "content": "m3"}) + trajectory._flush_json(force=True) + assert rendered == ["m0", "m1", "m2", "m3"] + + # Value replacement in place (as image eviction does) is detected. + history[1]["content"] = "edited" + trajectory._flush_json(force=True) + assert rendered[4:] == ["edited", "m2", "m3"] + # So is a rewrite that removes messages. + history[:] = [history[0], history[3]] + trajectory._flush_json(force=True) + assert [m["content"] for m in _messages(tmp_path / "trace.json")] == ["m0", "m3"] + assert rendered[7:] == ["m3"]