diff --git a/agent/moa_loop.py b/agent/moa_loop.py index e2b8c4e31c79..222506d8b17b 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -430,6 +430,7 @@ def _run_reference( temperature: float | None = None, max_tokens: int | None = None, reference_timeout: float | None = None, + context_length_cache: Any = None, ) -> tuple[str, str, Any]: """Call one reference model and return ``(label, text, accounting)``. @@ -465,6 +466,20 @@ def _run_reference( # trimmed view (_reference_messages) already strips the agent's own # system prompt, so this is the only system message the reference sees. messages = [{"role": "system", "content": _REFERENCE_SYSTEM_PROMPT}, *ref_messages] + # Trim to fit THIS reference model's context window. Reference models + # may have a smaller window than the aggregator (e.g. kimi-k2.7-code + # @ 262K advising a glm-5.2 @ 1M conversation); without this trim the + # provider returns a hard HTTP 400 which the except below silently + # converts to a [failed: …] note (issue #60345). Estimated AFTER the + # advisory system prompt is prepended so its tokens count against the + # budget too. + messages = _trim_messages_for_reference( + messages, + slot, + runtime, + reserve_output_tokens=max_tokens, + context_length_cache=context_length_cache, + ) # Apply the same Anthropic-style prompt-caching decoration the main # agent loop applies (system_and_3 breakpoints). The advisory view is # append-only across iterations (new turns append before the trailing @@ -567,6 +582,144 @@ def _run_reference( ) +# Output-token headroom reserved inside the reference's context window when +# the preset does not cap advisor output (reference_max_tokens=None). Roughly +# one long-form advisory answer; generous enough for thinking models' visible +# output without starving the input budget. +_REFERENCE_DEFAULT_OUTPUT_RESERVE = 8192 + +# Additional estimation slack: estimate_messages_tokens_rough is a rough +# chars/4 heuristic and providers tokenize less favorably on code/JSON-heavy +# transcripts, so keep a safety fraction of the window unbudgeted. +_REFERENCE_TRIM_SAFETY_FRACTION = 0.10 + + +def _trim_messages_for_reference( + messages: list[dict[str, Any]], + slot: dict[str, str], + runtime: dict[str, Any], + *, + reserve_output_tokens: int | None = None, + context_length_cache: Any = None, +) -> list[dict[str, Any]]: + """Trim an advisory request to fit within a reference model's context window. + + Reference models may have a smaller context window than the aggregator or + the main conversation. Without this trim, a reference whose window is + exceeded gets a hard HTTP 400 from the provider, which ``_run_reference``'s + try/except silently converts to a ``[failed: …]`` note — the MoA turn + silently degrades to fewer references (issue #60345). + + ``messages`` is the FULL request as it will be sent — the advisory system + prompt already prepended — so the estimate covers everything the provider + will count. The budget reserves ``reserve_output_tokens`` (the preset's + ``reference_max_tokens`` when set, else a sane constant) for the model's + response plus a safety fraction for estimator error. + + Trimming drops the OLDEST conversation frames (right after the system + prompt) and preserves two invariants of the advisory view, which is + text-only user/assistant turns (``_reference_messages`` renders tool + calls/results inline, so there are no tool-result frames to orphan): + + - the system prompt (index 0) is always kept; + - the first non-system message stays ``user``-first — after each pop, + any now-leading assistant turns are popped too, so no provider ever + sees an assistant-first conversation; + - the trailing user turn (the synthetic judge-the-state marker) and at + least one preceding turn are always kept, even if still over budget — + a too-long-but-recent view beats an empty request. + + ``context_length_cache`` is an optional per-turn dict keyed by + ``(provider, model)`` so one fan-out (and every iteration reusing the + cache) resolves each model's window at most once instead of re-probing + metadata sources per-reference-per-iteration. When the window cannot be + resolved, messages are returned unchanged. + """ + if not messages: + return messages + + from agent.model_metadata import ( + estimate_messages_tokens_rough, + get_model_context_length, + ) + + model = str(slot.get("model") or "") + provider = str(runtime.get("provider") or slot.get("provider") or "") + if not model: + return messages + + cache_key = (provider, model) + context_length: int | None = None + if isinstance(context_length_cache, dict) and cache_key in context_length_cache: + context_length = context_length_cache[cache_key] + else: + try: + context_length = get_model_context_length( + model=model, + base_url=str(runtime.get("base_url") or ""), + api_key=str(runtime.get("api_key") or ""), + provider=provider, + ) + except Exception: + logger.debug( + "MoA reference context-length resolution failed for %s", + _slot_label(slot), + ) + context_length = None + if isinstance(context_length_cache, dict): + # Cache failures too (as None) — a flaky metadata source should + # not be re-probed for every reference of every iteration. + context_length_cache[cache_key] = context_length + + if not isinstance(context_length, int) or context_length <= 0: + return messages + + reserve = ( + int(reserve_output_tokens) + if isinstance(reserve_output_tokens, int) and reserve_output_tokens > 0 + else _REFERENCE_DEFAULT_OUTPUT_RESERVE + ) + budget = int(context_length * (1.0 - _REFERENCE_TRIM_SAFETY_FRACTION)) - reserve + if budget <= 0: + return messages + + estimated = estimate_messages_tokens_rough(messages) + if estimated <= budget: + return messages + + has_system = bool(messages) and messages[0].get("role") == "system" + head = [messages[0]] if has_system else [] + body = list(messages[1:] if has_system else messages) + + # Keep the trailing user turn plus at least one preceding turn. + while len(body) > 2 and estimate_messages_tokens_rough(head + body) > budget: + body.pop(0) + # Preserve the user-first invariant: never leave the advisory + # conversation starting on an assistant turn after a pop. + while len(body) > 2 and body[0].get("role") == "assistant": + body.pop(0) + # The loop can stop with two frames left where the first is an + # assistant turn — enforce user-first even then (a lone trailing user + # turn is a valid request; an assistant-first one is not). + while len(body) > 1 and body[0].get("role") == "assistant": + body.pop(0) + + trimmed = head + body + dropped = len(messages) - len(trimmed) + if dropped: + logger.info( + "MoA reference %s: estimated %d tokens exceeds budget %d " + "(window %d, output reserve %d); dropped %d oldest message(s).", + _slot_label(slot), + estimated, + budget, + context_length, + reserve, + dropped, + ) + return trimmed + + _REFERENCE_POLL_INTERVAL_S = 5.0 # Sentinel text for a reference slot whose wait was aborted by a user @@ -637,6 +790,11 @@ def _run_references_parallel( completed = 0 executor = ThreadPoolExecutor(max_workers=workers) interrupted = False + # Per-fan-out context-length cache shared by every reference worker, so + # duplicate (provider, model) slots resolve their window once per turn + # instead of re-probing metadata sources per reference (dict get/set is + # GIL-atomic; a rare duplicate probe on a first-use race is harmless). + _ctx_len_cache: dict[tuple[str, str], int | None] = {} try: for idx, slot in enumerate(reference_models): if slot.get("provider") == "moa": @@ -654,6 +812,7 @@ def _run_references_parallel( temperature=temperature, max_tokens=max_tokens, reference_timeout=reference_timeout, + context_length_cache=_ctx_len_cache, ) ] = idx diff --git a/tests/run_agent/test_moa_loop_mode.py b/tests/run_agent/test_moa_loop_mode.py index e1f03c27a2d2..6e259eb9838c 100644 --- a/tests/run_agent/test_moa_loop_mode.py +++ b/tests/run_agent/test_moa_loop_mode.py @@ -2478,3 +2478,167 @@ moa: usage2, cost2 = facade.consume_reference_usage() assert usage2.input_tokens == 0 assert cost2 is None + + +class _CountingCtxLen: + """Stub for get_model_context_length that counts resolutions.""" + + def __init__(self, value): + self.value = value + self.calls = 0 + + def __call__(self, **kwargs): + self.calls += 1 + if isinstance(self.value, Exception): + raise self.value + return self.value + + +def _trim(messages, *, window=1000, reserve=None, cache=None, counting=None, + monkeypatch=None): + from agent import model_metadata, moa_loop + + stub = counting or _CountingCtxLen(window) + monkeypatch.setattr(model_metadata, "get_model_context_length", stub) + return moa_loop._trim_messages_for_reference( + messages, + {"provider": "openrouter", "model": "small-window"}, + {"provider": "openrouter", "model": "small-window"}, + reserve_output_tokens=reserve, + context_length_cache=cache, + ) + + +def _advisory_view(n_pairs, chunk="x" * 400): + """A text-only advisory view: system + n user/assistant pairs + trailing user.""" + msgs = [{"role": "system", "content": "advisory system prompt"}] + for i in range(n_pairs): + msgs.append({"role": "user", "content": f"u{i} {chunk}"}) + msgs.append({"role": "assistant", "content": f"a{i} {chunk}"}) + msgs.append({"role": "user", "content": "judge the state above"}) + return msgs + + +def test_reference_trim_untouched_when_within_window(monkeypatch): + msgs = _advisory_view(2) + out = _trim(list(msgs), window=10_000_000, monkeypatch=monkeypatch) + assert out == msgs + + +def test_reference_trim_drops_oldest_and_keeps_invariants(monkeypatch): + from agent.model_metadata import estimate_messages_tokens_rough + + msgs = _advisory_view(30) + out = _trim(list(msgs), window=4000, reserve=100, monkeypatch=monkeypatch) + + # Something was dropped and it fits the (window*0.9 - reserve) budget… + assert len(out) < len(msgs) + # System prompt survives at index 0. + assert out[0]["role"] == "system" + # User-first after the system prompt — never assistant-first. + assert out[1]["role"] == "user" + # Trailing synthetic user turn survives. + assert out[-1] == msgs[-1] + # Oldest frames were the ones dropped: the kept body is a contiguous + # suffix of the original body. + assert out[1:] == msgs[len(msgs) - len(out) + 1:] + + +def test_reference_trim_estimates_after_system_prompt(monkeypatch): + """The advisory system prompt counts against the budget: a view that fits + without it but not with it must still be trimmed.""" + from agent import model_metadata, moa_loop + + msgs = _advisory_view(6) + body_tokens = model_metadata.estimate_messages_tokens_rough(msgs[1:]) + total_tokens = model_metadata.estimate_messages_tokens_rough(msgs) + assert total_tokens > body_tokens + + # Pick a window where (body fits) but (system+body does not): + # budget = window*0.9 - reserve(100) must sit between the two. + window = int((body_tokens + (total_tokens - body_tokens) / 2 + 100) / 0.9) + out = _trim(list(msgs), window=window, reserve=100, monkeypatch=monkeypatch) + assert len(out) < len(msgs) + assert out[0]["role"] == "system" + + +def test_reference_trim_reserves_output_tokens(monkeypatch): + """With a huge output reserve the budget shrinks and forces a trim that + a reserve-less estimate would not need.""" + msgs = _advisory_view(10) + from agent.model_metadata import estimate_messages_tokens_rough + + total = estimate_messages_tokens_rough(msgs) + window = int((total + 500) / 0.9) # fits easily with a small reserve + out_small = _trim(list(msgs), window=window, reserve=100, monkeypatch=monkeypatch) + assert out_small == msgs + out_big = _trim(list(msgs), window=window, reserve=total, monkeypatch=monkeypatch) + assert len(out_big) < len(msgs) + + +def test_reference_trim_unresolvable_window_is_a_noop(monkeypatch): + msgs = _advisory_view(3) + stub = _CountingCtxLen(RuntimeError("metadata down")) + out = _trim(list(msgs), counting=stub, monkeypatch=monkeypatch) + assert out == msgs + + +def test_reference_trim_context_length_cache_hits_once(monkeypatch): + """A shared per-turn cache resolves each (provider, model) window once.""" + cache = {} + stub = _CountingCtxLen(10_000_000) + msgs = _advisory_view(2) + for _ in range(4): + _trim(list(msgs), cache=cache, counting=stub, monkeypatch=monkeypatch) + assert stub.calls == 1 + assert cache == {("openrouter", "small-window"): 10_000_000} + + +def test_reference_trim_caches_resolution_failures(monkeypatch): + """A failing metadata source is probed once, not per reference call.""" + cache = {} + stub = _CountingCtxLen(RuntimeError("metadata down")) + msgs = _advisory_view(2) + for _ in range(3): + out = _trim(list(msgs), cache=cache, counting=stub, monkeypatch=monkeypatch) + assert out == msgs + assert stub.calls == 1 + assert cache == {("openrouter", "small-window"): None} + + +def test_run_reference_trims_oversized_view_before_calling(monkeypatch): + """End-to-end: _run_reference sends a trimmed request for a small-window + model instead of letting the provider 400.""" + from agent import model_metadata, moa_loop + + sent = {} + + def fake_call_llm(**kwargs): + sent.update(kwargs) + return _response_with_usage("advice") + + monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm) + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda slot: {"provider": "openrouter", "model": slot["model"]}, + ) + monkeypatch.setattr( + model_metadata, "get_model_context_length", lambda **k: 3000 + ) + + view = _advisory_view(40)[1:] # advisory view has no system prompt + label, text, _acct = moa_loop._run_reference( + {"provider": "openrouter", "model": "small-window"}, + view, + max_tokens=200, + ) + + assert text == "advice" + sent_messages = sent["messages"] + # System prompt was prepended and survived the trim. + assert sent_messages[0]["role"] == "system" + # The request actually shrank. + assert len(sent_messages) < len(view) + 1 + # And ends with the synthetic trailing user turn. + assert sent_messages[-1]["role"] == "user"