From 68cd7557311f7d0d1a16e9b69082fde5cdda0c9e Mon Sep 17 00:00:00 2001 From: srojk34 <286497132+srojk34@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:27:36 +0300 Subject: [PATCH] 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. --- agent/conversation_loop.py | 9 ++- agent/moa_loop.py | 111 +++++++++++++++++++++----- tests/run_agent/test_moa_loop_mode.py | 76 ++++++++++++++++++ 3 files changed, 177 insertions(+), 19 deletions(-) diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index dec6a43cfd2..a764501feba 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -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): diff --git a/agent/moa_loop.py b/agent/moa_loop.py index aaa6654246b..5b72add97dc 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -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, ) diff --git a/tests/run_agent/test_moa_loop_mode.py b/tests/run_agent/test_moa_loop_mode.py index 49db70e2968..aa1db3d8936 100644 --- a/tests/run_agent/test_moa_loop_mode.py +++ b/tests/run_agent/test_moa_loop_mode.py @@ -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(