diff --git a/agent/moa_loop.py b/agent/moa_loop.py index 5b72add97dc..e2b8c4e31c7 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -11,6 +11,7 @@ from __future__ import annotations import hashlib import logging import re +import threading from concurrent.futures import ThreadPoolExecutor, wait as _futures_wait from typing import Any @@ -568,6 +569,11 @@ def _run_reference( _REFERENCE_POLL_INTERVAL_S = 5.0 +# Sentinel text for a reference slot whose wait was aborted by a user +# interrupt. Shared by _run_references_parallel (which writes it) and the +# facade cache logic (which must never cache it as real advice). +_INTERRUPTED_REFERENCE_NOTE = "[skipped: interrupted by user]" + def _run_references_parallel( reference_models: list[dict[str, Any]], @@ -578,6 +584,7 @@ def _run_references_parallel( progress_callback: Any = None, reference_timeout: float | None = None, agent: Any = None, + late_accounting_sink: Any = None, ) -> list[tuple[str, str, Any]]: """Fan out all reference models in parallel, returning outputs in order. @@ -678,23 +685,44 @@ def _run_references_parallel( if results[idx] is not None: continue if future.cancel(): + # Never dispatched — genuinely nothing was billed. results[idx] = ( _slot_label(reference_models[idx]), - "[skipped: interrupted by user]", + _INTERRUPTED_REFERENCE_NOTE, _RefAccounting(CanonicalUsage()), ) elif future.done(): - # Finished between the interrupt check and now. + # Finished between the interrupt check and now — the call + # completed and billed, so keep its REAL output and + # accounting rather than zeroing it with a placeholder. results[idx] = future.result() else: # Already running — cannot be force-killed (see # docstring); leave it be so the caller isn't blocked, - # and note that its output was abandoned. + # and note that its output was abandoned. The provider + # call is still in flight and WILL bill when it + # completes, so hand its eventual accounting to the + # caller's sink instead of silently dropping it. + label = _slot_label(reference_models[idx]) results[idx] = ( - _slot_label(reference_models[idx]), - "[skipped: interrupted by user]", + label, + _INTERRUPTED_REFERENCE_NOTE, _RefAccounting(CanonicalUsage()), ) + if late_accounting_sink is not None: + def _record_late(f: Any, _label: str = label) -> None: + try: + _lbl, _txt, _acct = f.result() + except Exception: # pragma: no cover - defensive + return + try: + late_accounting_sink(_label, _acct) + except Exception: # pragma: no cover - defensive + logger.debug( + "MoA: late accounting sink failed for %s", + _label, + ) + future.add_done_callback(_record_late) finally: executor.shutdown(wait=not interrupted, cancel_futures=interrupted) @@ -1183,6 +1211,10 @@ class MoAChatCompletions: self._pending_reference_usage: Any = CanonicalUsage() self._pending_reference_cost: Any = None + # Guards pending usage/cost against concurrent late-accounting + # callbacks (see _record_late_reference_accounting), which fire on + # executor worker threads after an interrupted fan-out returns. + self._accounting_lock = threading.Lock() # Resolved aggregator slot ({provider, model, ...}) from the most recent # create(); read by session cost accounting to price the aggregator's # acting turn at its real model instead of the virtual preset name. @@ -1215,12 +1247,42 @@ class MoAChatCompletions: """ from agent.usage_pricing import CanonicalUsage - usage = self._pending_reference_usage or CanonicalUsage() - cost = self._pending_reference_cost - self._pending_reference_usage = CanonicalUsage() - self._pending_reference_cost = None + with self._accounting_lock: + usage = self._pending_reference_usage or CanonicalUsage() + cost = self._pending_reference_cost + self._pending_reference_usage = CanonicalUsage() + self._pending_reference_cost = None return usage, cost + def _record_late_reference_accounting(self, label: str, accounting: Any) -> None: + """Fold a late-completing interrupted reference's real spend in. + + When a user interrupt aborts the fan-out wait, references already in + flight keep running (they cannot be force-killed) and DO bill when + they complete. Their placeholder results carry zeroed accounting, so + without this hook that spend would vanish from session accounting. + The fan-out registers this as a done-callback on abandoned futures; + it folds the eventual real usage/cost into the pending totals, where + the next ``consume_reference_usage`` pick-up records it. Thread-safe: + done-callbacks fire on executor worker threads. + """ + from agent.usage_pricing import CanonicalUsage + + if not isinstance(accounting, _RefAccounting): + return + with self._accounting_lock: + if isinstance(accounting.usage, CanonicalUsage): + self._pending_reference_usage = ( + self._pending_reference_usage or CanonicalUsage() + ) + accounting.usage + if accounting.cost_usd is not None: + self._pending_reference_cost = ( + self._pending_reference_cost or 0 + ) + accounting.cost_usd + logger.debug( + "MoA: recorded late accounting for interrupted reference %s", label + ) + def consume_and_save_trace( self, session_id: Any = None, aggregator_output_fallback: Any = None ) -> None: @@ -1565,9 +1627,11 @@ class MoAChatCompletions: # References already ran (and were accounted) earlier this turn; # this create() is a repeat tool-iteration reusing the cached # advice. Charging their tokens/cost again here would multiply - # advisor spend by the tool-iteration count, so pending is zero. - self._pending_reference_usage = CanonicalUsage() - self._pending_reference_cost = None + # advisor spend by the tool-iteration count, so nothing new is + # deposited — but do NOT zero the pending totals: a + # late-completing interrupted reference may have deposited its + # real spend since the last consume(), and that must survive + # until the next consume_reference_usage() pick-up. # Likewise no trace on a cache HIT — the full turn was already # traced on the MISS that ran the references. A repeat iteration is # not a new MoA turn. @@ -1594,9 +1658,23 @@ class MoAChatCompletions: progress_callback=_progress, reference_timeout=reference_timeout, agent=self._agent, + late_accounting_sink=self._record_late_reference_accounting, ) - self._ref_cache_key = _cache_key - self._ref_cache_outputs = list(reference_outputs) + interrupted_any = any( + text == _INTERRUPTED_REFERENCE_NOTE + for _lbl, text, _acct in reference_outputs + ) + if interrupted_any: + # An interrupted fan-out is a partial snapshot, not real + # advice for this state. Caching it would replay the + # placeholder notes on every subsequent iteration of the + # turn (a cache HIT never re-runs the references), so leave + # the cache empty and let the next create() re-run them. + self._ref_cache_key = None + self._ref_cache_outputs = [] + else: + self._ref_cache_key = _cache_key + self._ref_cache_outputs = list(reference_outputs) # Sum the advisor fan-out's token usage AND cost so the caller can # fold advisor spend into session accounting exactly once per turn. # Only the freshly run references (cache MISS) contribute; a cache @@ -1613,8 +1691,17 @@ class MoAChatCompletions: _ref_usage = _ref_usage + _acct.usage if _acct.cost_usd is not None: _ref_cost = (_ref_cost or 0) + _acct.cost_usd - self._pending_reference_usage = _ref_usage - self._pending_reference_cost = _ref_cost + with self._accounting_lock: + # Fold (don't overwrite): a late-completing interrupted + # reference from a PREVIOUS turn may have deposited its real + # spend here between consume() calls — keep it. + self._pending_reference_usage = ( + self._pending_reference_usage or CanonicalUsage() + ) + _ref_usage + if _ref_cost is not None: + self._pending_reference_cost = ( + self._pending_reference_cost or 0 + ) + _ref_cost # Stash the full reference fan-out for trace persistence. The # aggregator input/label are filled in below once agg_messages is # built; the aggregator OUTPUT is stitched in by the caller diff --git a/tests/run_agent/test_moa_loop_mode.py b/tests/run_agent/test_moa_loop_mode.py index aa1db3d8936..e1f03c27a2d 100644 --- a/tests/run_agent/test_moa_loop_mode.py +++ b/tests/run_agent/test_moa_loop_mode.py @@ -2310,3 +2310,171 @@ def test_moa_facade_acts_aggregator_alone_when_all_references_fail_silent( assert "[failed:" not in prompt assert "Reference models unavailable" not in prompt assert "Mixture of Agents reference context" not in prompt + + +def test_interrupted_but_completed_reference_keeps_real_accounting(monkeypatch): + """A reference that finishes between the interrupt check and the reap + must keep its REAL output and accounting — the call billed.""" + from concurrent.futures import wait as real_wait + + from agent import moa_loop + + monkeypatch.setattr(moa_loop, "get_transport", lambda *_a, **_k: None) + monkeypatch.setattr(moa_loop, "_REFERENCE_POLL_INTERVAL_S", 0.05) + + fake_agent = SimpleNamespace(_interrupt_requested=True) + + def fake_call_llm(**kwargs): + return _response_with_usage("slowish output", prompt=11, completion=4) + + # Force the exact race: the wait loop reports the future as still + # pending (so the interrupt path is taken) even though the underlying + # call has already completed — the reap must then hit the done() branch + # and keep the real result instead of writing a placeholder. + def fake_wait(pending, timeout=None): + real_wait(pending) # let the call actually finish (it billed) + return set(), set(pending) # report it as still pending + + monkeypatch.setattr(moa_loop, "_futures_wait", fake_wait) + monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm) + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda slot: {"provider": slot["provider"], "model": slot["model"]}, + ) + + out = moa_loop._run_references_parallel( + [{"provider": "slowish", "model": "m1"}], + [{"role": "user", "content": "hi"}], + agent=fake_agent, + ) + + # The completed call's real output + usage must survive the reap. + assert out[0][1] == "slowish output" + acct = out[0][2] + assert isinstance(acct, moa_loop._RefAccounting) + assert acct.usage.input_tokens == 11 + + +def test_late_completing_interrupted_reference_feeds_accounting_sink(monkeypatch): + """A reference still in flight at interrupt time gets a placeholder in + the results, but its eventual REAL accounting must reach the sink.""" + import threading + import time + + from agent import moa_loop + + monkeypatch.setattr(moa_loop, "get_transport", lambda *_a, **_k: None) + monkeypatch.setattr(moa_loop, "_REFERENCE_POLL_INTERVAL_S", 0.05) + + fake_agent = SimpleNamespace(_interrupt_requested=False) + release = threading.Event() + sink_calls = [] + sink_seen = threading.Event() + + def sink(label, accounting): + sink_calls.append((label, accounting)) + sink_seen.set() + + def fake_call_llm(**kwargs): + if kwargs["provider"] == "fast": + fake_agent._interrupt_requested = True + return _response("fast output") + # wedged: blocks past the interrupt, completes later. + release.wait(timeout=5) + return _response_with_usage("late output", prompt=21, completion=2) + + monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm) + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda slot: {"provider": slot["provider"], "model": slot["model"]}, + ) + + out = moa_loop._run_references_parallel( + [ + {"provider": "fast", "model": "m1"}, + {"provider": "wedged", "model": "m2"}, + ], + [{"role": "user", "content": "hi"}], + agent=fake_agent, + late_accounting_sink=sink, + ) + + # The wedged slot returned a placeholder with zeroed accounting… + assert out[1][1] == moa_loop._INTERRUPTED_REFERENCE_NOTE + assert out[1][2].usage.input_tokens == 0 + + # …then completes late; its real billed usage must reach the sink. + release.set() + assert sink_seen.wait(timeout=5), "late accounting sink never called" + label, acct = sink_calls[0] + assert "wedged" in label + assert acct.usage.input_tokens == 21 + + +def test_facade_does_not_cache_interrupted_reference_results(monkeypatch, tmp_path): + """An interrupted fan-out is a partial snapshot — caching it would replay + placeholder notes on every later iteration of the turn. The facade must + leave the cache empty so the next create() re-runs the references, and + a late-completing reference's real spend must land in pending usage.""" + from agent import moa_loop + from agent.usage_pricing import CanonicalUsage + + home = tmp_path / ".hermes" + home.mkdir() + (home / "config.yaml").write_text( + """ +moa: + default_preset: review + presets: + review: + reference_models: + - provider: openrouter + model: advisor + aggregator: + provider: openrouter + model: aggregator +""".strip(), + encoding="utf-8", + ) + monkeypatch.setenv("HERMES_HOME", str(home)) + + interrupted_outputs = [ + ( + "openrouter:advisor", + moa_loop._INTERRUPTED_REFERENCE_NOTE, + moa_loop._RefAccounting(CanonicalUsage()), + ) + ] + + def fake_fanout(*args, **kwargs): + return list(interrupted_outputs) + + monkeypatch.setattr(moa_loop, "_run_references_parallel", fake_fanout) + monkeypatch.setattr(moa_loop, "call_llm", lambda **k: _response("acted")) + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda slot: {"provider": slot["provider"], "model": slot["model"]}, + ) + + facade = moa_loop.MoAChatCompletions("review") + facade.create(messages=[{"role": "user", "content": "go"}], tools=[]) + + # Interrupted results must not be cached as this state's advice. + assert facade._ref_cache_key is None + assert facade._ref_cache_outputs == [] + + # A late completion depositing real spend is picked up by consume(). + facade._record_late_reference_accounting( + "openrouter:advisor", + moa_loop._RefAccounting(CanonicalUsage(input_tokens=33), 0.42), + ) + usage, cost = facade.consume_reference_usage() + assert usage.input_tokens == 33 + assert cost == pytest.approx(0.42) + # And consume() drained it — no double count. + usage2, cost2 = facade.consume_reference_usage() + assert usage2.input_tokens == 0 + assert cost2 is None