From 5de65624d13f61211ed89e7a3a2805f9bae88c18 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 1 Jul 2026 04:07:46 -0700 Subject: [PATCH] fix(moa): capture streamed aggregator output into full-turn traces (#56312) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MoA full-turn traces (moa.save_traces) recorded the aggregator's acting output only on the non-streaming path, where it's captured inline at call time. On the streaming path — which every hermes chat --query run and every live gateway/CLI turn takes — the aggregator's raw token stream is handed to the live consumer, so the trace left output=null and only pointed at the session-db assistant row. An offline audit of a benchmark run (HermesBench drives --query) then couldn't see what the aggregator produced without hand-joining to state.db. Capture the resolved streamed acting text at trace-flush time (the agent already holds it in _current_streamed_assistant_text) and fold it into the trace, so the record is self-contained in both modes. New output_location value inline_from_stream marks a streamed turn whose text was captured this way; a genuinely empty acting turn (pure tool call) still points at the session db, matching state.db exactly. Touches only the trace side-channel — no change to the acting path, message history, role alternation, or prompt cache. - agent/moa_loop.py: consume_and_save_trace(..., aggregator_output_fallback) on both the facade and the MoAClient wrapper; prefer inline capture, fall back to the resolved streamed text. - agent/moa_trace.py: embed the fallback; add inline_from_stream location. - agent/conversation_loop.py: pass _current_streamed_assistant_text at flush. - tests: 5 cases across streaming / non-streaming / empty-fallback / no-double-write. --- agent/conversation_loop.py | 14 +- agent/moa_loop.py | 30 +++- agent/moa_trace.py | 34 ++-- .../agent/test_moa_trace_streamed_capture.py | 155 ++++++++++++++++++ 4 files changed, 217 insertions(+), 16 deletions(-) create mode 100644 tests/agent/test_moa_trace_streamed_capture.py diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index a60af85c8e6..3e253c82a53 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -1967,10 +1967,20 @@ def run_conversation( # Flush the full-turn MoA trace (references + aggregator I/O) # to disk when moa.save_traces is on. No-op otherwise and # for non-MoA clients. Uses the live session_id so traces - # land in the right per-session file. + # land in the right per-session file. On the streaming path + # the aggregator's output wasn't captured inline (its raw + # token stream went to the live consumer), so pass the + # resolved streamed acting text as a fallback — makes the + # trace self-contained instead of only pointing at state.db. if _moa_client is not None and hasattr(_moa_client, "consume_and_save_trace"): try: - _moa_client.consume_and_save_trace(agent.session_id) + _agg_streamed_text = ( + getattr(agent, "_current_streamed_assistant_text", "") or "" + ) + _moa_client.consume_and_save_trace( + agent.session_id, + aggregator_output_fallback=_agg_streamed_text or None, + ) except Exception as _moa_trace_exc: # pragma: no cover - defensive logger.debug("MoA trace flush failed: %s", _moa_trace_exc) prompt_tokens = canonical_usage.prompt_tokens diff --git a/agent/moa_loop.py b/agent/moa_loop.py index 015bc23ac00..891d256f673 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -640,13 +640,24 @@ class MoAChatCompletions: self._pending_reference_cost = None return usage, cost - def consume_and_save_trace(self, session_id: Any = None) -> None: + def consume_and_save_trace( + self, session_id: Any = None, aggregator_output_fallback: Any = None + ) -> None: """Flush the pending full-turn trace to disk, if one is pending. No-op when tracing is off (``save_moa_turn`` checks the config), when there is no pending trace (a cache-HIT iteration ran no references), or when the aggregator input was never recorded. Clears the pending trace so a repeat consume cannot double-write. Best-effort — never raises. + + ``aggregator_output_fallback`` is the aggregator's resolved acting text + as the caller already holds it in memory (the streamed assistant text). + On the streaming path the aggregator's output could not be captured + inline at ``create()`` time (the raw token stream was handed to the live + consumer), so ``pending["aggregator_output"]`` is None; we fold the + caller's resolved text in here so the trace is self-contained in BOTH + streaming and non-streaming modes. Non-streaming already has the inline + output and ignores the fallback. """ pending = self._pending_trace self._pending_trace = None @@ -656,6 +667,11 @@ class MoAChatCompletions: from agent.moa_trace import save_moa_turn agg_slot = pending.get("aggregator_slot") or {} + # Prefer the inline capture (non-streaming); fall back to the + # caller's resolved streamed text when streaming left it None. + agg_output = pending.get("aggregator_output") + if agg_output is None and aggregator_output_fallback: + agg_output = aggregator_output_fallback save_moa_turn( session_id=session_id, preset_name=pending.get("preset", ""), @@ -665,7 +681,7 @@ class MoAChatCompletions: aggregator_provider=agg_slot.get("provider"), aggregator_temperature=pending.get("aggregator_temperature"), aggregator_input_messages=pending.get("aggregator_input_messages"), - aggregator_output=pending.get("aggregator_output"), + aggregator_output=agg_output, aggregator_streamed=bool(pending.get("aggregator_streamed")), ) except Exception as exc: # pragma: no cover - tracing must never break a turn @@ -885,9 +901,15 @@ class MoAClient: """ return self.chat.completions.consume_reference_usage() - def consume_and_save_trace(self, session_id: Any = None) -> None: + def consume_and_save_trace( + self, session_id: Any = None, aggregator_output_fallback: Any = None + ) -> None: """Flush the pending full-turn MoA trace via the completions facade. No-op unless ``moa.save_traces`` is enabled and a turn is pending. + ``aggregator_output_fallback`` supplies the resolved acting text so the + streaming path's trace is self-contained (see the facade docstring). """ - return self.chat.completions.consume_and_save_trace(session_id) + return self.chat.completions.consume_and_save_trace( + session_id, aggregator_output_fallback=aggregator_output_fallback + ) diff --git a/agent/moa_trace.py b/agent/moa_trace.py index a18a26df86e..37a51700812 100644 --- a/agent/moa_trace.py +++ b/agent/moa_trace.py @@ -112,11 +112,13 @@ def save_moa_turn( Best-effort: any failure is logged at debug and swallowed — tracing must never break a live turn. Called once per turn on a reference cache MISS. - ``aggregator_output`` is the aggregator's synthesized text when it was - captured inline (non-streaming path — the eval / quiet-mode path). When the - aggregator streamed to a live consumer, ``aggregator_streamed`` is True and - the output is delivered as the turn's assistant message in the session - store instead; the trace records the full aggregator INPUT either way. + ``aggregator_output`` is the aggregator's synthesized text. On the + non-streaming path (eval / quiet-mode / subagents) it was captured inline + at call time. On the streaming path it is captured after the fact from the + caller's resolved assistant text (``aggregator_output_fallback`` in + ``consume_and_save_trace``) so the trace is self-contained either way; if + that resolved text was unavailable, it falls back to None and the record + points at the session store via ``output_location``. """ base = _traces_enabled_and_dir() if base is None: @@ -124,6 +126,16 @@ def save_moa_turn( try: base.mkdir(parents=True, exist_ok=True) path = base / f"{_sanitize_session_id(session_id)}.jsonl" + # output_location tells an offline reader where the acting text lives: + # embedded here when we have it (both non-streaming inline capture and + # streaming after-the-fact capture), else the session-db assistant row. + _have_output = bool(aggregator_output) + if not aggregator_streamed: + _output_location = "inline" + elif _have_output: + _output_location = "inline_from_stream" + else: + _output_location = "assistant_message_in_session_db" record = { "ts": time.time(), "session_id": session_id, @@ -140,11 +152,13 @@ def save_moa_turn( "input_messages": aggregator_input_messages, "output": aggregator_output, "streamed": aggregator_streamed, - # When streamed, the aggregator's acting output is persisted as - # the turn's assistant message in state.db (see the session - # store); it is not duplicated here. - "output_location": "assistant_message_in_session_db" - if aggregator_streamed else "inline", + # Where the aggregator's acting output lives for this record. + # "inline" — non-streaming inline capture + # "inline_from_stream" — streamed, then captured from the + # caller's resolved assistant text + # "assistant_message_in_session_db" — streamed and the resolved + # text was unavailable at flush time + "output_location": _output_location, }, } with path.open("a", encoding="utf-8") as f: diff --git a/tests/agent/test_moa_trace_streamed_capture.py b/tests/agent/test_moa_trace_streamed_capture.py new file mode 100644 index 00000000000..b149182c992 --- /dev/null +++ b/tests/agent/test_moa_trace_streamed_capture.py @@ -0,0 +1,155 @@ +"""Tests for MoA trace aggregator-output capture across streaming modes. + +The MoA full-turn trace (opt-in ``moa.save_traces``) must record the +aggregator's acting output whether the aggregator ran non-streaming (inline +capture at call time) or streaming (captured after the fact from the caller's +resolved assistant text). Before the streamed-capture fix, a streamed +aggregator left ``output: null`` in the trace and only pointed at state.db, +so an offline audit of a benchmark run (which drives the streaming display +path via ``hermes chat --query``) couldn't see what the aggregator actually +produced without joining to the session DB by hand. + +These exercise the real ``consume_and_save_trace`` → ``save_moa_turn`` path +with real file I/O against a temp HERMES_HOME — no mocks on the write path. +""" + +from __future__ import annotations + +import json + +import pytest + +from agent.moa_loop import MoAChatCompletions + + +def _enable_traces(tmp_path, monkeypatch): + """Point HERMES_HOME at a temp dir and turn moa.save_traces on.""" + hermes_home = tmp_path / ".hermes" + hermes_home.mkdir() + monkeypatch.setenv("HERMES_HOME", str(hermes_home)) + + # save_moa_turn reads config via hermes_cli.config.load_config; stub it to + # return traces-on so the test doesn't depend on a real config file. + import agent.moa_trace as moa_trace + + monkeypatch.setattr( + moa_trace, + "load_config", + lambda: {"moa": {"save_traces": True}}, + raising=False, + ) + # load_config is imported lazily inside _traces_enabled_and_dir; patch the + # source module attribute it imports from as well. + import hermes_cli.config as cfg + + monkeypatch.setattr( + cfg, "load_config", lambda: {"moa": {"save_traces": True}}, raising=False + ) + return hermes_home / "moa-traces" + + +def _make_completions_with_pending(streamed: bool, inline_output): + """Build a MoAChatCompletions with a pending trace mimicking one turn.""" + mc = MoAChatCompletions.__new__(MoAChatCompletions) + mc._pending_trace = { + "preset": "closed", + "reference_outputs": [], # references not under test here + "aggregator_label": "openrouter:anthropic/claude-opus-4.8", + "aggregator_slot": { + "model": "anthropic/claude-opus-4.8", + "provider": "openrouter", + }, + "aggregator_temperature": 0.4, + "aggregator_input_messages": [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "do the thing"}, + ], + "aggregator_output": inline_output, + "aggregator_streamed": streamed, + } + return mc + + +def _read_single_trace(trace_dir, session_id): + path = trace_dir / f"{session_id}.jsonl" + assert path.exists(), f"trace file not written: {path}" + lines = path.read_text().strip().split("\n") + assert len(lines) == 1 + return json.loads(lines[0]) + + +def test_streamed_aggregator_output_captured_from_fallback(tmp_path, monkeypatch): + """Streaming turn: inline output is None, fallback text is embedded.""" + trace_dir = _enable_traces(tmp_path, monkeypatch) + mc = _make_completions_with_pending(streamed=True, inline_output=None) + + mc.consume_and_save_trace( + "sess_streamed", + aggregator_output_fallback="the acting aggregator answer", + ) + + rec = _read_single_trace(trace_dir, "sess_streamed") + agg = rec["aggregator"] + assert agg["streamed"] is True + assert agg["output"] == "the acting aggregator answer" + assert agg["output_location"] == "inline_from_stream" + + +def test_non_streaming_prefers_inline_over_fallback(tmp_path, monkeypatch): + """Non-streaming turn keeps its inline capture even if a fallback is passed.""" + trace_dir = _enable_traces(tmp_path, monkeypatch) + mc = _make_completions_with_pending( + streamed=False, inline_output="inline captured text" + ) + + mc.consume_and_save_trace( + "sess_inline", + aggregator_output_fallback="SHOULD NOT BE USED", + ) + + rec = _read_single_trace(trace_dir, "sess_inline") + agg = rec["aggregator"] + assert agg["streamed"] is False + assert agg["output"] == "inline captured text" + assert agg["output_location"] == "inline" + + +def test_streamed_without_fallback_points_to_session_db(tmp_path, monkeypatch): + """Streaming turn with no resolvable text falls back to the state.db pointer.""" + trace_dir = _enable_traces(tmp_path, monkeypatch) + mc = _make_completions_with_pending(streamed=True, inline_output=None) + + mc.consume_and_save_trace("sess_nofb", aggregator_output_fallback=None) + + rec = _read_single_trace(trace_dir, "sess_nofb") + agg = rec["aggregator"] + assert agg["streamed"] is True + assert agg["output"] is None + assert agg["output_location"] == "assistant_message_in_session_db" + + +def test_pending_trace_cleared_after_flush(tmp_path, monkeypatch): + """A second flush is a no-op (pending cleared) — never double-writes.""" + trace_dir = _enable_traces(tmp_path, monkeypatch) + mc = _make_completions_with_pending(streamed=True, inline_output=None) + + mc.consume_and_save_trace("sess_once", aggregator_output_fallback="x") + # Second call: pending is None now, must not append a second line. + mc.consume_and_save_trace("sess_once", aggregator_output_fallback="y") + + path = trace_dir / "sess_once.jsonl" + lines = path.read_text().strip().split("\n") + assert len(lines) == 1 + + +def test_empty_fallback_string_treated_as_missing(tmp_path, monkeypatch): + """An empty-string fallback must not override to '' — treated as absent.""" + trace_dir = _enable_traces(tmp_path, monkeypatch) + mc = _make_completions_with_pending(streamed=True, inline_output=None) + + mc.consume_and_save_trace("sess_empty", aggregator_output_fallback="") + + rec = _read_single_trace(trace_dir, "sess_empty") + agg = rec["aggregator"] + assert agg["output"] is None + assert agg["output_location"] == "assistant_message_in_session_db"