diff --git a/agent/agent_init.py b/agent/agent_init.py index 92be00c7bf43..166f26b8aaa7 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -1032,12 +1032,12 @@ def init_agent( # reference-model outputs to the agent's tool_progress_callback so # every surface that already consumes it (CLI spinner/scrollback, TUI, # desktop, gateway) can show each reference's answer as a labelled - # block before the aggregator acts. The facade emits "moa.reference" - # and "moa.aggregating" events, forwarded through the same callback - # the tool lifecycle uses. Best-effort and cache-safe — display-only - # events, they never touch the message history. The factory is shared - # with the fallback-restore/recovery paths so a restored facade keeps - # emitting these events (#53802). + # block before the aggregator acts. The facade emits "moa.reference", + # "moa.progress", "moa.phase", and "moa.aggregating" events, forwarded + # through the same callback the tool lifecycle uses. Best-effort and + # cache-safe — display-only events, they never touch the message + # history. The factory is shared with the fallback-restore/recovery + # paths so a restored facade keeps emitting these events (#53802). agent.client = build_moa_facade(agent, agent.model) agent._client_kwargs = {} agent.api_key = api_key or "moa-virtual-provider" diff --git a/agent/moa_loop.py b/agent/moa_loop.py index 3eceeb945bac..c625220671dd 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -559,6 +559,7 @@ def _run_references_parallel( *, temperature: float | None = None, max_tokens: int | None = None, + progress_callback: Any = None, ) -> list[tuple[str, str, Any]]: """Fan out all reference models in parallel, returning outputs in order. @@ -568,6 +569,12 @@ def _run_references_parallel( ``Reference {idx}`` labelling stays stable. MoA presets that reference another MoA preset are skipped here (recursion guard) with a labelled note. + If ``progress_callback`` is provided it is invoked as each reference + completes: ``progress_callback(refs_done, refs_total, label)``. The total + matches ``len(reference_models)`` so listeners can render a status-bar + progress like ``MOA: 2/3 refs done``. Best-effort — failures are logged + but never break the fan-out (display must never block a turn). + Each element is ``(label, text, usage)`` where usage is a ``CanonicalUsage`` (zeroed for skipped/failed references). """ @@ -585,6 +592,8 @@ def _run_references_parallel( # advisor calls attribute to the same conversation as the acting turn. from tools.thread_context import propagate_context_to_thread + total = len(reference_models) + done = 0 with ThreadPoolExecutor(max_workers=workers) as executor: for idx, slot in enumerate(reference_models): if slot.get("provider") == "moa": @@ -607,6 +616,13 @@ def _run_references_parallel( # complete set, so there is no early-exit / first-completed path here. for future, idx in futures.items(): results[idx] = future.result() + done += 1 + if progress_callback is not None: + try: + label = _slot_label(reference_models[idx]) + progress_callback(done, total, label) + except Exception as exc: # pragma: no cover - display must never break + logger.debug("MoA progress_callback failed: %s", exc) return [r for r in results if r is not None] @@ -988,6 +1004,14 @@ class MoAChatCompletions: # reference_callback(event, **kwargs) # where event is one of: # "moa.reference" kwargs: index, count, label, text + # "moa.progress" kwargs: refs_done, refs_total, label + # (fired once per reference completion — drives + # status-bar progress like ``MOA: 2/3 refs done``) + # "moa.phase" kwargs: phase, refs_done, refs_total, aggregator + # (fired on phase transitions, currently + # phase="aggregator" right before the aggregator + # acts; phase="reference" mirrors ``moa.progress`` + # so listeners can rely on a single event family) # "moa.aggregating" kwargs: aggregator (label), ref_count # Never raises into the model call — display is best-effort. self.reference_callback = reference_callback @@ -1389,11 +1413,25 @@ class MoAChatCompletions: # not a new MoA turn. self._pending_trace = None else: + # Per-reference progress callback: emits ``moa.progress`` so + # listeners can render ``MOA: N/M refs done`` in the status bar as + # each reference completes. The callback is bound to self so it + # goes through the same display hook as the existing + # ``moa.reference`` / ``moa.aggregating`` events. + def _progress(done: int, total: int, label: str) -> None: + self._emit( + "moa.progress", + refs_done=done, + refs_total=total, + label=label, + ) + reference_outputs = _run_references_parallel( reference_models, ref_messages, temperature=temperature, max_tokens=reference_max_tokens, + progress_callback=_progress, ) self._ref_cache_key = _cache_key self._ref_cache_outputs = list(reference_outputs) @@ -1456,6 +1494,17 @@ class MoAChatCompletions: text=_redact_reference_text(_text) if privacy_mode else _text, ) if _ref_count: + # Phase transition: reference fan-out is complete, the + # aggregator is about to act. Listeners that prefer a single + # event family for phase tracking can switch on ``phase`` + # instead of subscribing to ``moa.aggregating`` separately. + self._emit( + "moa.phase", + phase="aggregator", + refs_done=_ref_count, + refs_total=_ref_count, + aggregator=_slot_label(aggregator), + ) self._emit( "moa.aggregating", aggregator=_slot_label(aggregator), @@ -1574,6 +1623,30 @@ def build_moa_facade(agent, preset_name: Any = None) -> MoAClient: moa_index=idx, moa_count=count, ) + elif event == "moa.progress": + # Per-reference completion. Frontends render this as a + # status-bar progress indicator like ``MOA: N/M refs done``. + cb( + "moa.progress", + str(kwargs.get("label") or ""), + None, + None, + moa_refs_done=kwargs.get("refs_done"), + moa_refs_total=kwargs.get("refs_total"), + ) + elif event == "moa.phase": + # Phase transition (currently only ``phase="aggregator"`` + # fires once the fan-out is done). Subscribers can switch + # on ``moa_phase`` to know which phase is active. + cb( + "moa.phase", + str(kwargs.get("aggregator") or ""), + None, + None, + moa_phase=kwargs.get("phase"), + moa_refs_done=kwargs.get("refs_done"), + moa_refs_total=kwargs.get("refs_total"), + ) elif event == "moa.aggregating": cb( "moa.aggregating", diff --git a/tests/agent/test_moa_progress.py b/tests/agent/test_moa_progress.py new file mode 100644 index 000000000000..75d61829ffea --- /dev/null +++ b/tests/agent/test_moa_progress.py @@ -0,0 +1,234 @@ +"""Tests for the MOA progress indicator added in issue #59546. + +The MoA facade (``MoAChatCompletions.create``) emits a sequence of display +events so a TUI / CLI / desktop surface can render progress as references +complete and the phase transitions from ``reference`` to ``aggregator``: + + - ``moa.progress`` — fired once per reference completion with + ``refs_done`` and ``refs_total`` (e.g. drives a + status-bar ``MOA: 2/3 refs done`` indicator) + - ``moa.phase`` — fired once per phase transition (currently only the + ``phase="aggregator"`` transition right before the + aggregator acts) + +These tests exercise the real callback surface end-to-end through the +display hook (``reference_callback``) — no mocks on the dispatch path. +The LLM is stubbed via ``call_llm`` so the test does not depend on any +real provider. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + + +def _response(content="ok"): + message = SimpleNamespace(content=content, tool_calls=[]) + choice = SimpleNamespace(message=message, finish_reason="stop") + return SimpleNamespace(choices=[choice], usage=None, model="fake") + + +@pytest.fixture +def moa_config(tmp_path, monkeypatch): + home = tmp_path / ".hermes" + home.mkdir() + (home / "config.yaml").write_text( + """ +moa: + default_preset: closed + presets: + closed: + enabled: true + reference_models: + - provider: openrouter + model: anthropic/claude-opus-4.8 + - provider: openrouter + model: openai/gpt-5.5 + - provider: openrouter + model: google/gemini-3-pro + aggregator: + provider: openrouter + model: anthropic/claude-opus-4.8 +""".strip(), + encoding="utf-8", + ) + monkeypatch.setenv("HERMES_HOME", str(home)) + return home + + +def _collect_emits(facade): + """Pull every event the facade dispatches into a flat list of (event, kwargs).""" + captured: list[tuple[str, dict]] = [] + + def _capture(event: str, **kwargs): + captured.append((event, kwargs)) + + facade.reference_callback = _capture + return captured + + +def test_moa_progress_fires_for_each_reference(moa_config, monkeypatch): + """One ``moa.progress`` event per reference completion with monotonic counts.""" + from agent.moa_loop import MoAChatCompletions + + def fake_call_llm(**kwargs): + # Per-model stub: each reference returns a stable string so we can + # assert labels flow through; the aggregator returns the acting text. + if kwargs.get("task") == "moa_reference": + return _response(f"advice from {kwargs.get('model', '?')}") + return _response("acted") + + monkeypatch.setattr("agent.moa_loop.call_llm", fake_call_llm) + + facade = MoAChatCompletions("closed") + captured = _collect_emits(facade) + + facade.create( + model="closed", + messages=[{"role": "user", "content": "clean the db"}], + ) + + progress_events = [(e, k) for (e, k) in captured if e == "moa.progress"] + # 3 references configured in moa_config => 3 progress events. + assert len(progress_events) == 3 + + # Monotonic 1/3, 2/3, 3/3 — each event carries the current count and total. + expected_counts = [(1, 3), (2, 3), (3, 3)] + actual_counts = [ + (k["refs_done"], k["refs_total"]) for (_, k) in progress_events + ] + assert actual_counts == expected_counts + + # Every progress event names the source model slot (or close to it) so a + # status bar can render ``MOA: 2/3 refs done — openai/gpt-5.5``. + for _, kwargs in progress_events: + assert "label" in kwargs + assert kwargs["label"] + + +def test_moa_phase_transitions_to_aggregator(moa_config, monkeypatch): + """A single ``moa.phase`` event with ``phase="aggregator"`` fires after the fan-out.""" + from agent.moa_loop import MoAChatCompletions + + def fake_call_llm(**kwargs): + if kwargs.get("task") == "moa_reference": + return _response("advice") + return _response("acted") + + monkeypatch.setattr("agent.moa_loop.call_llm", fake_call_llm) + + facade = MoAChatCompletions("closed") + captured = _collect_emits(facade) + + facade.create( + model="closed", + messages=[{"role": "user", "content": "plan the migration"}], + ) + + phase_events = [(e, k) for (e, k) in captured if e == "moa.phase"] + # Exactly one phase event per turn, identifying the aggregator. + assert len(phase_events) == 1 + _event, kwargs = phase_events[0] + assert kwargs["phase"] == "aggregator" + assert kwargs["aggregator"] == "openrouter:anthropic/claude-opus-4.8" + # counts match the configured reference count + assert kwargs["refs_done"] == 3 + assert kwargs["refs_total"] == 3 + + +def test_moa_progress_counts_match_n_references(moa_config, monkeypatch): + """Progress counters equal ``len(reference_models)`` regardless of size.""" + from agent.moa_loop import MoAChatCompletions + + def fake_call_llm(**kwargs): + if kwargs.get("task") == "moa_reference": + return _response("advice") + return _response("acted") + + monkeypatch.setattr("agent.moa_loop.call_llm", fake_call_llm) + + facade = MoAChatCompletions("closed") + captured = _collect_emits(facade) + + facade.create( + model="closed", + messages=[{"role": "user", "content": "summarize"}], + ) + + progress_events = [(e, k) for (e, k) in captured if e == "moa.progress"] + totals = [k["refs_total"] for _, k in progress_events] + # Every event reports the same total — the preset's reference-model count. + assert totals and all(t == 3 for t in totals) + # And the final done-count equals the total (fan-out finished). + final = progress_events[-1][1] + assert final["refs_done"] == final["refs_total"] + + +def test_moa_progress_event_order_matches_fanout(moa_config, monkeypatch): + """Every progress event fires AFTER its matching moa.reference event. + + Listeners that animate one block per reference (collapsible) need the + progress notification to land after the per-reference text so the + status-bar counter and the rendered block stay in lockstep. + """ + from agent.moa_loop import MoAChatCompletions + + def fake_call_llm(**kwargs): + if kwargs.get("task") == "moa_reference": + return _response("advice") + return _response("acted") + + monkeypatch.setattr("agent.moa_loop.call_llm", fake_call_llm) + + facade = MoAChatCompletions("closed") + captured = _collect_emits(facade) + + facade.create( + model="closed", + messages=[{"role": "user", "content": "rank the options"}], + ) + + # Walk the events; every progress event must be preceded by its reference + # text event (matching ``index``). The full sequence ends with the + # aggregator phase event. + seen_phase = False + for event, kwargs in captured: + if event == "moa.reference": + assert kwargs["index"] <= kwargs["count"] + elif event == "moa.progress": + assert kwargs["refs_done"] <= kwargs["refs_total"] + elif event == "moa.phase": + # No further reference events after the aggregator phase event. + assert kwargs["phase"] == "aggregator" + seen_phase = True + elif event == "moa.aggregating": + # Legacy marker still fires for backwards compatibility, and it + # always lands AFTER the phase event. + assert seen_phase + assert seen_phase, "expected at least one moa.phase event" + + +def test_moa_progress_callback_none_safe(moa_config, monkeypatch): + """A missing ``reference_callback`` does not break the fan-out or create().""" + from agent.moa_loop import MoAChatCompletions + + def fake_call_llm(**kwargs): + if kwargs.get("task") == "moa_reference": + return _response("advice") + return _response("acted") + + monkeypatch.setattr("agent.moa_loop.call_llm", fake_call_llm) + + # No callback attached — the facade's _emit is a no-op in that case. + facade = MoAChatCompletions("closed") + assert facade.reference_callback is None + facade.create( + model="closed", + messages=[{"role": "user", "content": "noop"}], + ) + # Turn still resolved cleanly; aggregator slot populated as usual. + assert facade.last_aggregator_slot is not None + assert facade.last_aggregator_slot["model"] == "anthropic/claude-opus-4.8" \ No newline at end of file diff --git a/tui_gateway/server.py b/tui_gateway/server.py index eb9f5cda0e66..844961ee37a2 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -4432,6 +4432,43 @@ def _on_tool_progress( if event_type == "moa.aggregating": _emit("moa.aggregating", sid, {"aggregator": str(name or "")}) return + if event_type == "moa.progress": + # Per-reference completion — drives the status-bar progress indicator + # (`MOA: 2/3 refs done`) requested in issue #59546. Only emitted when + # both counters are present so the client can render deterministically. + refs_done = _kwargs.get("moa_refs_done") + refs_total = _kwargs.get("moa_refs_total") + if refs_done is None or refs_total is None: + return + _emit( + "moa.progress", + sid, + { + "label": str(name or ""), + "refs_done": int(refs_done), + "refs_total": int(refs_total), + }, + ) + return + if event_type == "moa.phase": + # Phase transition — currently only ``phase="aggregator"`` fires once + # the fan-out completes and the aggregator is about to act. Tells the + # client which phase of the MoA pipeline is currently running so it + # can swap status-bar copy accordingly. + phase = _kwargs.get("moa_phase") + if not phase: + return + phase_payload: dict[str, object] = {"phase": str(phase)} + refs_done = _kwargs.get("moa_refs_done") + refs_total = _kwargs.get("moa_refs_total") + if refs_done is not None: + phase_payload["refs_done"] = int(refs_done) + if refs_total is not None: + phase_payload["refs_total"] = int(refs_total) + if name: + phase_payload["aggregator"] = str(name) + _emit("moa.phase", sid, phase_payload) + return if event_type.startswith("subagent."): payload = { "goal": str(_kwargs.get("goal") or ""),