fix(moa): allow a user interrupt to abort the reference fan-out wait

agent/tool_executor.py's concurrent tool batch checks agent._interrupt_requested
and aborts the wait early; agent/moa_loop.py's _run_references_parallel had
no equivalent, so a MoA-enabled turn blocked on ThreadPoolExecutor.result()
until every reference model finished or hit its own individual
auxiliary.moa_reference timeout -- there was no way for the user to abort a
live turn mid-fanout.

Thread an optional `agent` parameter through aggregate_moa_context ->
_run_references_parallel (used when MoA references run alongside the main
model) and MoAClient/MoAChatCompletions (used when the MoA preset itself is
the acting model), then poll concurrent.futures.wait() in
_REFERENCE_POLL_INTERVAL_S slices instead of blocking on future.result() per
reference, checking agent._interrupt_requested each cycle.

Deliberately scoped to interrupt/cancel only -- no new or changed timeout
value, so this doesn't overlap open PRs #53784/#53875 (which lower the
per-reference timeout default but don't add interrupt support). `agent` is
optional and defaults to None, so any caller that doesn't pass it keeps
today's uninterruptible blocking behavior unchanged.
This commit is contained in:
srojk34 2026-07-01 14:27:36 +03:00 committed by Teknium
parent 62c2b299a3
commit 68cd755731
3 changed files with 177 additions and 19 deletions

View file

@ -1171,10 +1171,17 @@ def run_conversation(
temperature=_preset_temperature(moa_config, "reference_temperature"),
aggregator_temperature=_preset_temperature(moa_config, "aggregator_temperature"),
reference_max_tokens=moa_config.get("reference_max_tokens"),
reference_timeout=float(moa_config.get("reference_timeout") or 30.0),
# None = no per-preset override; inherit
# auxiliary.moa_reference.timeout via call_llm.
reference_timeout=(
float(moa_config["reference_timeout"])
if moa_config.get("reference_timeout")
else None
),
degraded_reference_policy=str(
moa_config.get("degraded_reference_policy") or "loud"
),
agent=agent,
)
if _moa_context:
for _msg in reversed(api_messages):

View file

@ -11,7 +11,7 @@ from __future__ import annotations
import hashlib
import logging
import re
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import ThreadPoolExecutor, wait as _futures_wait
from typing import Any
from agent.auxiliary_client import call_llm
@ -566,6 +566,9 @@ def _run_reference(
)
_REFERENCE_POLL_INTERVAL_S = 5.0
def _run_references_parallel(
reference_models: list[dict[str, Any]],
ref_messages: list[dict[str, Any]],
@ -574,6 +577,7 @@ def _run_references_parallel(
max_tokens: int | None = None,
progress_callback: Any = None,
reference_timeout: float | None = None,
agent: Any = None,
) -> list[tuple[str, str, Any]]:
"""Fan out all reference models in parallel, returning outputs in order.
@ -590,7 +594,23 @@ def _run_references_parallel(
but never break the fan-out (display must never block a turn).
Each element is ``(label, text, accounting)`` where accounting is a
``_RefAccounting`` object (zeroed for skipped/failed references).
``_RefAccounting`` object (zeroed for skipped/failed/interrupted
references).
When *agent* is given, the fan-out is interruptible: waiting for the
batch is broken into ``_REFERENCE_POLL_INTERVAL_S``-second polls (instead
of one blocking ``future.result()`` per reference) so a user interrupt
mid-turn can abort the wait mirroring the same interrupt check
``agent.tool_executor`` already applies to its own concurrent tool
batch. This does not add or change any per-reference *timeout* (that is
``reference_timeout`` / ``auxiliary.moa_reference.timeout``, resolved
elsewhere) it only lets the caller stop waiting early. References
already in flight cannot be forcibly killed (``call_llm`` is a blocking
HTTP call with no interrupt hook of its own, same limitation
tool_executor has for tools without an interrupt check); an interrupted
reference's own timeout still reaps its thread independently. *agent* is
optional and defaults to ``None``, preserving the uninterruptible
blocking behavior for any caller that doesn't pass it.
"""
from agent.usage_pricing import CanonicalUsage
@ -598,7 +618,7 @@ def _run_references_parallel(
return []
results: list[tuple[str, str, Any] | None] = [None] * len(reference_models)
futures = {}
futures: dict[Any, int] = {}
workers = min(_MAX_REFERENCE_WORKERS, len(reference_models))
# Reference slots run on bare executor threads, which start with an empty
# contextvars.Context — propagate the parent turn's context (approval
@ -607,8 +627,10 @@ def _run_references_parallel(
from tools.thread_context import propagate_context_to_thread
total = len(reference_models)
done = 0
with ThreadPoolExecutor(max_workers=workers) as executor:
completed = 0
executor = ThreadPoolExecutor(max_workers=workers)
interrupted = False
try:
for idx, slot in enumerate(reference_models):
if slot.get("provider") == "moa":
results[idx] = (
@ -627,17 +649,54 @@ def _run_references_parallel(
reference_timeout=reference_timeout,
)
] = idx
# Collect every reference before returning — the aggregator needs the
# 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)
# complete set, so there is no early-exit / first-completed path
# here, other than a user interrupt. Progress callbacks fire as each
# reference completes so frontends can render "MOA: k/n refs done".
pending = set(futures)
while pending:
done, pending = _futures_wait(pending, timeout=_REFERENCE_POLL_INTERVAL_S)
for future in done:
idx = futures[future]
results[idx] = future.result()
completed += 1
if progress_callback is not None:
try:
label = _slot_label(reference_models[idx])
progress_callback(completed, total, label)
except Exception as exc: # pragma: no cover - display must never break
logger.debug("MoA progress_callback failed: %s", exc)
if not pending:
break
if agent is not None and getattr(agent, "_interrupt_requested", False):
interrupted = True
break
if interrupted:
for future, idx in futures.items():
if results[idx] is not None:
continue
if future.cancel():
results[idx] = (
_slot_label(reference_models[idx]),
"[skipped: interrupted by user]",
_RefAccounting(CanonicalUsage()),
)
elif future.done():
# Finished between the interrupt check and now.
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.
results[idx] = (
_slot_label(reference_models[idx]),
"[skipped: interrupted by user]",
_RefAccounting(CanonicalUsage()),
)
finally:
executor.shutdown(wait=not interrupted, cancel_futures=interrupted)
return [r for r in results if r is not None]
@ -908,6 +967,7 @@ def aggregate_moa_context(
reference_max_tokens: int | None = None,
reference_timeout: float | None = None,
degraded_reference_policy: str = "loud",
agent: Any = None,
) -> str:
"""Run configured reference models and synthesize their advice.
@ -927,6 +987,9 @@ def aggregate_moa_context(
like ``reference_max_tokens``, ``call_llm`` omits temperature when None
so the provider default applies matching single-model agent behavior.
Presets may still pin explicit values.
``agent``, when passed, lets the reference fan-out be aborted early on a
user interrupt see ``_run_references_parallel``'s docstring.
"""
reference_models = [slot for slot in reference_models if slot.get("enabled", True)]
reference_outputs: list[tuple[str, str, Any]] = []
@ -937,6 +1000,7 @@ def aggregate_moa_context(
temperature=temperature,
max_tokens=reference_max_tokens,
reference_timeout=reference_timeout,
agent=agent,
)
successful_outputs = _successful_references(reference_outputs)
@ -1074,7 +1138,7 @@ def _attach_reference_guidance(agg_messages: list[dict[str, Any]], guidance: str
class MoAChatCompletions:
"""OpenAI-chat-compatible facade where the aggregator is the acting model."""
def __init__(self, preset_name: str, reference_callback: Any = None):
def __init__(self, preset_name: str, reference_callback: Any = None, agent: Any = None):
self.preset_name = preset_name or "default"
# Optional display hook. Called as reference outputs become available so
# frontends can show each reference model's answer as a labelled block
@ -1093,6 +1157,11 @@ class MoAChatCompletions:
# "moa.aggregating" kwargs: aggregator (label), ref_count
# Never raises into the model call — display is best-effort.
self.reference_callback = reference_callback
# Back-reference to the owning AIAgent, so the reference fan-out can
# check agent._interrupt_requested (see _run_references_parallel).
# Optional — a caller that doesn't pass it just keeps the fan-out
# uninterruptible, as it was before.
self._agent = agent
# State-scoped reference cache. The agent loop calls create() once per
# tool-loop iteration; references should re-run whenever the task STATE
# advances — i.e. on every new user message AND every new tool result —
@ -1524,6 +1593,7 @@ class MoAChatCompletions:
max_tokens=reference_max_tokens,
progress_callback=_progress,
reference_timeout=reference_timeout,
agent=self._agent,
)
self._ref_cache_key = _cache_key
self._ref_cache_outputs = list(reference_outputs)
@ -1675,9 +1745,11 @@ class MoAChatCompletions:
class MoAClient:
def __init__(self, preset_name: str, reference_callback: Any = None):
def __init__(self, preset_name: str, reference_callback: Any = None, agent: Any = None):
self.chat = type("_MoAChat", (), {})()
self.chat.completions = MoAChatCompletions(preset_name, reference_callback=reference_callback)
self.chat.completions = MoAChatCompletions(
preset_name, reference_callback=reference_callback, agent=agent,
)
def consume_reference_usage(self) -> Any:
"""Pop the pending reference-fan-out usage from the completions facade.
@ -1785,4 +1857,7 @@ def build_moa_facade(agent, preset_name: Any = None) -> MoAClient:
return MoAClient(
str(preset_name or getattr(agent, "model", None) or "default"),
reference_callback=_moa_reference_relay,
# Thread the agent through so the reference fan-out wait can be
# aborted on a user interrupt (see _run_references_parallel).
agent=agent,
)

View file

@ -1101,6 +1101,82 @@ def test_references_run_in_parallel(monkeypatch):
assert out[0][1] == "resp-p1"
def test_references_parallel_without_agent_is_unaffected(monkeypatch):
"""No agent passed (the pre-fix call shape) must behave exactly as
before: block until every reference completes, no interrupt check."""
import time
from agent import moa_loop
monkeypatch.setattr(moa_loop, "get_transport", lambda *_a, **_k: None)
# Poll interval shorter than the reference's own sleep so the assertion
# below would catch a regression that waits a whole poll cycle extra.
monkeypatch.setattr(moa_loop, "_REFERENCE_POLL_INTERVAL_S", 0.05)
def slow_call_llm(**kwargs):
time.sleep(0.2)
return _response(f"resp-{kwargs['provider']}")
monkeypatch.setattr(moa_loop, "call_llm", slow_call_llm)
refs = [{"provider": "p1", "model": "ok"}]
out = moa_loop._run_references_parallel(
refs, [{"role": "user", "content": "hi"}],
)
assert out[0][1] == "resp-p1"
def test_references_parallel_interrupt_aborts_wait(monkeypatch):
"""A user interrupt mid-fanout must stop the wait instead of blocking
until every reference (including a wedged one) finishes or times out on
its own mirroring the interrupt check agent.tool_executor already
applies to its own concurrent tool batch."""
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_wedged = threading.Event()
def fake_call_llm(**kwargs):
if kwargs["provider"] == "fast":
# Simulate the interrupt arriving right after the fast reference
# finishes, while the wedged one is still in flight.
fake_agent._interrupt_requested = True
return _response("fast output")
# "wedged" — never returns within the test unless released, standing
# in for a reference whose own (possibly very long) timeout hasn't
# elapsed yet.
release_wedged.wait(timeout=5)
return _response("should not be observed")
monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm)
refs = [
{"provider": "fast", "model": "m1"},
{"provider": "wedged", "model": "m2"},
]
try:
start = time.monotonic()
out = moa_loop._run_references_parallel(
refs, [{"role": "user", "content": "hi"}], agent=fake_agent,
)
elapsed = time.monotonic() - start
# Must return promptly once interrupted, not block for the wedged
# reference's full (5s test-simulated) duration.
assert elapsed < 2.0, f"interrupt did not abort the wait (took {elapsed:.2f}s)"
assert out[0][1] == "fast output"
assert "interrupted" in out[1][1]
finally:
release_wedged.set() # don't leak a blocked thread past the test
def _ref_config(home):
home.mkdir()
(home / "config.yaml").write_text(