mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-31 19:16:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

fix(moa): keep real accounting for interrupted-but-billed references; don't cache interrupted results

Browse source 
Follow-ups for salvaged #56344:

- A reference that completes between the interrupt check and the reap
  keeps its REAL output and accounting (the provider call billed) instead
  of being zeroed with a placeholder.
- A reference still in flight at interrupt time gets a placeholder in the
  results, but its future now carries a done-callback that folds the
  eventual real usage/cost into the facade's pending accounting
  (late_accounting_sink -> _record_late_reference_accounting), so billed
  spend is never silently dropped. Pending totals are folded (not
  overwritten) and guarded by a lock since done-callbacks fire on
  executor worker threads.
- Interrupted placeholder results are no longer written into the facade's
  turn-scoped reference cache: a cache HIT never re-runs references, so
  caching a partial snapshot would replay '[skipped: interrupted by
  user]' notes for the rest of the turn. The cache is left empty and the
  next create() re-runs the fan-out.
This commit is contained in:
Teknium 2026-07-23 12:28:51 -07:00
parent 68cd755731
commit 55f3826224
2 changed files with 271 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

119
agent/moa_loop.py
View file

@ -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

168
tests/run_agent/test_moa_loop_mode.py
View file

@ -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