diff --git a/agent/moa_loop.py b/agent/moa_loop.py index ccaebda8f07..9700f4abe85 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -173,18 +173,19 @@ def _slot_runtime(slot: dict[str, str]) -> dict[str, Any]: return out -def _maybe_apply_advisor_cache_control( +def _maybe_apply_moa_cache_control( messages: list[dict[str, Any]], runtime: dict[str, Any], ) -> list[dict[str, Any]]: - """Decorate an advisor request with cache_control when its route honors it. + """Decorate an advisor or aggregator request with cache_control when its + route honors it. Reuses the SAME policy function as the main agent loop - (``anthropic_prompt_cache_policy``) resolved against the advisor slot's - own provider/base_url/api_mode/model, and the SAME breakpoint layout - (``apply_anthropic_cache_control``, system_and_3). This keeps advisor - calls decorated exactly like an acting agent on that provider would be — - no MoA-specific caching logic to drift. + (``anthropic_prompt_cache_policy``) resolved against the slot's own + provider/base_url/api_mode/model, and the SAME breakpoint layout + (``apply_anthropic_cache_control``, system_and_3). This keeps advisor and + aggregator calls decorated exactly like an acting agent on that provider + would be — no MoA-specific caching logic to drift. Returns the messages unchanged on any resolution error or when the policy says the route doesn't honor markers. @@ -196,8 +197,8 @@ def _maybe_apply_advisor_cache_control( from agent.prompt_caching import apply_anthropic_cache_control # The policy function reads agent.* only as fallbacks for kwargs we - # don't pass; provide a stub so an advisor slot is judged purely on - # its own resolved runtime. + # don't pass; provide a stub so the slot is judged purely on its own + # resolved runtime. stub = SimpleNamespace(provider="", base_url="", api_mode="", model="") should_cache, native_layout = anthropic_prompt_cache_policy( stub, @@ -212,7 +213,7 @@ def _maybe_apply_advisor_cache_control( messages, native_anthropic=native_layout ) except Exception as exc: # pragma: no cover - decoration must never break a call - logger.debug("advisor cache_control decoration skipped: %s", exc) + logger.debug("MoA cache_control decoration skipped: %s", exc) return messages @@ -268,7 +269,7 @@ def _run_reference( # caching is opt-in per request. OpenAI-family advisors are untouched # (their caching is automatic; markers are ignored harmlessly, but we # only decorate when the policy says the route honors them). - messages = _maybe_apply_advisor_cache_control(messages, runtime) + messages = _maybe_apply_moa_cache_control(messages, runtime) response = call_llm( task="moa_reference", messages=messages, @@ -617,13 +618,27 @@ def aggregate_moa_context( ) agg_label = _slot_label(aggregator) + agg_runtime = _slot_runtime(aggregator) try: + # Same cache_control decoration as _run_reference's advisor calls + # (see _maybe_apply_moa_cache_control) — this synthesis call is a + # third, independent MoA call path that 22c5048d9 did not cover (it + # only restored caching for the acting-aggregator turn in the + # persistent `provider: moa` model and for advisor fan-out). Without + # it, the one-shot `/moa ` command's synthesis call re-bills + # its full input (system-less prompt containing every joined + # reference output) on every invocation with zero cache_control + # breakpoints, even when the resolved aggregator slot is a + # cache-honoring route (e.g. Claude on OpenRouter/native Anthropic). + agg_messages = _maybe_apply_moa_cache_control( + [{"role": "user", "content": synth_prompt}], agg_runtime + ) response = call_llm( task="moa_aggregator", - messages=[{"role": "user", "content": synth_prompt}], + messages=agg_messages, temperature=aggregator_temperature, max_tokens=max_tokens, - **_slot_runtime(aggregator), + **agg_runtime, ) synthesis = _extract_text(response) except Exception as exc: diff --git a/tests/agent/test_moa_aggregator_cache_control.py b/tests/agent/test_moa_aggregator_cache_control.py new file mode 100644 index 00000000000..0c7baf441f1 --- /dev/null +++ b/tests/agent/test_moa_aggregator_cache_control.py @@ -0,0 +1,114 @@ +"""Regression test: the MoA aggregator's one-shot synthesis call +(``aggregate_moa_context``, used by the ``/moa `` command) must get +the same Anthropic-style prompt-caching decoration as the acting-aggregator +turn (``MoAChatCompletions.create``) and the advisor fan-out +(``_run_reference``). + +22c5048d9 ("fix(moa): restore prompt caching for the aggregator and +advisors") fixed the other two MoA call paths but never touched +``aggregate_moa_context`` — a third, independent call path with its own +``call_llm(task="moa_aggregator", ...)`` invocation. Without this fix, every +``/moa `` one-shot call re-bills its full input (system-less prompt +containing all joined reference outputs) with zero cache_control breakpoints, +even when the resolved aggregator slot is a cache-honoring route. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + + +def _response(content="synthesized guidance"): + message = SimpleNamespace(content=content, tool_calls=[]) + choice = SimpleNamespace(message=message, finish_reason="stop") + return SimpleNamespace(choices=[choice], usage=None, model="fake") + + +@pytest.fixture +def captured_calls(monkeypatch): + calls = [] + + def fake_call_llm(**kwargs): + calls.append(kwargs) + return _response() + + monkeypatch.setattr("agent.moa_loop.call_llm", fake_call_llm) + monkeypatch.setattr( + "agent.moa_loop._run_references_parallel", + lambda *a, **k: [("advisor-a", "advice from a", None)], + ) + return calls + + +def _aggregator_kwargs(calls): + return next(c for c in calls if c.get("task") == "moa_aggregator") + + +def test_aggregator_synthesis_gets_cache_control_on_native_anthropic_route( + captured_calls, monkeypatch +): + """A cache-honoring aggregator slot (native Anthropic) must get + cache_control breakpoints on its synthesis call.""" + from agent import moa_loop + + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda slot: { + "provider": "anthropic", + "model": "claude-opus-4.8", + "base_url": "", + "api_mode": "anthropic_messages", + }, + ) + + moa_loop.aggregate_moa_context( + user_prompt="what should I do next?", + api_messages=[{"role": "user", "content": "help me plan"}], + reference_models=[{"provider": "openrouter", "model": "openai/gpt-5.5"}], + aggregator={"provider": "anthropic", "model": "claude-opus-4.8"}, + ) + + agg_kwargs = _aggregator_kwargs(captured_calls) + synth_message = agg_kwargs["messages"][0] + assert synth_message["role"] == "user" + content = synth_message["content"] + # Native Anthropic layout places cache_control on inner content blocks, + # so a cached message's content is a list of blocks rather than a bare + # string once decorated. + assert isinstance(content, list), "expected native cache_control block layout" + assert any( + isinstance(block, dict) and "cache_control" in block for block in content + ), "aggregator synthesis message must carry a cache_control breakpoint" + + +def test_aggregator_synthesis_untouched_on_non_caching_route( + captured_calls, monkeypatch +): + """A non-cache-honoring aggregator slot (plain OpenAI) must not be + decorated — proves the guard doesn't over-fire.""" + from agent import moa_loop + + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda slot: { + "provider": "openai", + "model": "gpt-5.5", + "base_url": "", + "api_mode": "chat_completions", + }, + ) + + moa_loop.aggregate_moa_context( + user_prompt="what should I do next?", + api_messages=[{"role": "user", "content": "help me plan"}], + reference_models=[{"provider": "openrouter", "model": "openai/gpt-5.5"}], + aggregator={"provider": "openai", "model": "gpt-5.5"}, + ) + + agg_kwargs = _aggregator_kwargs(captured_calls) + synth_message = agg_kwargs["messages"][0] + assert isinstance(synth_message["content"], str), "must stay undecorated (plain string content)"