fix(moa): apply prompt-caching decoration to the aggregator's one-shot synthesis call

22c5048d9 restored Anthropic-style cache_control for two of MoA's three
call paths: the acting aggregator (MoAChatCompletions.create, the
persistent `provider: moa` model) and the advisor fan-out (_run_reference).
aggregate_moa_context() -- the /moa <prompt> one-shot command's synthesis
call -- is the third, independent call path and was never covered: its
call_llm(task="moa_aggregator", ...) sent a single undecorated user message
containing the full joined reference output, re-billing the entire input on
every invocation even when the resolved aggregator slot is a cache-honoring
route (Claude on OpenRouter/native Anthropic, MiniMax, Qwen/DashScope).

- Generalize _maybe_apply_advisor_cache_control to
  _maybe_apply_moa_cache_control (it never had advisor-specific logic --
  same policy function, same breakpoint layout as the main loop, judged
  purely on the passed-in runtime) and reuse it in aggregate_moa_context
  the same way _run_reference already does.
- Compute _slot_runtime(aggregator) once and reuse it for both the
  decoration call and the call_llm kwargs, instead of calling it twice.

Mutation-verified: reverting the moa_loop.py change makes the new
regression test fail by asserting a plain string aggregator-message
content where the cache-honoring case expects native cache_control
content blocks.
This commit is contained in:
srojk34 2026-07-04 16:09:24 +03:00 committed by kshitij
parent 5daa5a0f2f
commit 2d3eac5fbd
2 changed files with 142 additions and 13 deletions

View file

@ -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 <prompt>` 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:

View file

@ -0,0 +1,114 @@
"""Regression test: the MoA aggregator's one-shot synthesis call
(``aggregate_moa_context``, used by the ``/moa <prompt>`` 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 <prompt>`` 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)"