From f3e62619c7f693f4747df5fceae2322c385ac4e3 Mon Sep 17 00:00:00 2001 From: teknium1 <127238744+teknium1@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:18:06 -0700 Subject: [PATCH] fix(moa): keep the aggregator prefix byte-stable across tool-loop iterations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Attach the per-turn reference guidance as its own trailing user message instead of merging it into a trailing user turn. Iteration 1 of every user turn ends on user(task), so the merged user(task + guidance) byte-differed from the user(task) that iteration 2+ replayed: the provider prefix cache collapsed to the system prompt + tool schemas on the second call of every turn (#112358 — cachedTokens pinned at the system prompt while uncached input grew 14k -> 214k on a Nous/OpenAI aggregator). peel_reference_guidance shrinks to the one attach shape (plain string, or its cache-decorated single text part) so a breakpoint never lands on the guidance. Converters that require strict alternation (Anthropic Messages, Converse, native Gemini) already merge adjacent same-role turns. Live wire probe (loopback server, two user turns x two iterations, real MoA facade -> call_llm(nous)): before, request i+1 diverged from request i at messages[1] (user task) on the second iteration of each turn; after, every request's messages[:-1] is a byte prefix of the next. The same probe shows the aggregator already carries the main loop's sticky session_id + tags on the nous provider, so the routing hypothesis is not a Hermes gap. --- agent/moa_loop.py | 45 ++++++----------- tests/agent/test_failover_identity.py | 3 +- .../test_moa_aggregator_prefix_stability.py | 49 +++++++++++++++++++ .../user-guide/features/mixture-of-agents.md | 2 +- 4 files changed, 68 insertions(+), 31 deletions(-) create mode 100644 tests/agent/test_moa_aggregator_prefix_stability.py diff --git a/agent/moa_loop.py b/agent/moa_loop.py index 1138e90deb..d24371cc16 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -889,26 +889,22 @@ def _completed_response_as_stream_chunk(response: Any) -> Any: def _attach_reference_guidance(agg_messages: list[dict[str, Any]], guidance: str) -> None: - """Attach the per-turn reference block at the END of the aggregator prompt. + """Attach the per-turn reference block as its OWN trailing user message. - The block varies per iteration; appending keeps ``[system][task][tool-history]`` - cache-stable. A trailing user turn is merged in place (string, or a new text part - AFTER the cache_control-marked part); otherwise a user message is appended (two - consecutive user turns would be rejected by strict providers). + The block varies per turn; appending keeps ``[system][task][tool-history]`` + cache-stable. It is never merged into a trailing user turn: iteration 1 of a + tool loop ends on ``user(task)``, and a merged ``user(task + guidance)`` byte-differs + from the ``user(task)`` every later iteration replays, so the provider prefix cache + collapsed to the system prompt on iteration 2 of every turn (#112358). Converters + that require strict alternation already merge adjacent same-role turns. """ - last = agg_messages[-1] if agg_messages else None - last_content = last.get("content") if last is not None and last.get("role") == "user" else None - if isinstance(last_content, str): - last["content"] = last_content + "\n\n" + guidance - elif isinstance(last_content, list): - last["content"] = [*last_content, {"type": "text", "text": "\n\n" + guidance}] - else: - agg_messages.append({"role": "user", "content": guidance}) + agg_messages.append({"role": "user", "content": guidance}) def peel_reference_guidance(messages: list[dict[str, Any]], guidance: Any) -> list[dict[str, Any]]: - """Exact inverse of ``_attach_reference_guidance`` (the three attach shapes), so a - cache breakpoint never lands on the turn-varying guidance. Inputs are not mutated.""" + """Exact inverse of ``_attach_reference_guidance`` (plain string, or its cache-decorated + single-text-part form), so a cache breakpoint never lands on the turn-varying guidance. + Inputs are not mutated.""" if not guidance or not messages: return messages guidance_text = str(guidance) @@ -916,21 +912,12 @@ def peel_reference_guidance(messages: list[dict[str, Any]], guidance: Any) -> li if not isinstance(last, dict) or last.get("role") != "user": return messages content = last.get("content") - if content == guidance_text: # shape (c): guidance was its own user message + if content == guidance_text: return list(messages[:-1]) - suffix = "\n\n" + guidance_text - if isinstance(content, str) and content.endswith(suffix): # shape (a): merged into a string turn - return [*messages[:-1], {**last, "content": content[: -len(suffix)]}] - if isinstance(content, list) and content: - last_part = content[-1] - if isinstance(last_part, dict) and last_part.get("type", "text") == "text": - text = last_part.get("text") or "" - if text in (suffix, guidance_text): - # Shape (b): guidance rode as its own trailing part. Guidance as the - # only content drops the whole message (mirrors shape c). - return list(messages[:-1]) if len(content) == 1 else [*messages[:-1], {**last, "content": list(content[:-1])}] - if text.endswith(suffix): - return [*messages[:-1], {**last, "content": [*content[:-1], {**last_part, "text": text[: -len(suffix)]}]}] + if isinstance(content, list) and len(content) == 1: + part = content[0] + if isinstance(part, dict) and part.get("type", "text") == "text" and (part.get("text") or "") == guidance_text: + return list(messages[:-1]) return messages diff --git a/tests/agent/test_failover_identity.py b/tests/agent/test_failover_identity.py index 2f4485a660..86e3ef0f20 100644 --- a/tests/agent/test_failover_identity.py +++ b/tests/agent/test_failover_identity.py @@ -271,7 +271,8 @@ class TestRedecoratePromptCacheOnPolicyChange: {"role": "system", "content": prompt}, {"role": "user", "content": "task"}, {"role": "assistant", "content": "ok"}, - {"role": "user", "content": "task\n\n" + guidance}, + {"role": "user", "content": "task"}, + {"role": "user", "content": guidance}, ] decorated = apply_anthropic_cache_control(base, native_anthropic=True) diff --git a/tests/agent/test_moa_aggregator_prefix_stability.py b/tests/agent/test_moa_aggregator_prefix_stability.py new file mode 100644 index 0000000000..bb336ff67b --- /dev/null +++ b/tests/agent/test_moa_aggregator_prefix_stability.py @@ -0,0 +1,49 @@ +"""MoA aggregator requests must grow as a byte-stable prefix across tool-loop iterations. + +Issue #112358: the reference guidance used to be merged INTO a trailing user turn on +iteration 1 of every user turn, so iteration 2's ``user(task)`` byte-differed from the one +the provider had just cached and the prompt cache collapsed to the system prompt. +""" + +from types import SimpleNamespace + +from agent import moa_loop + + +def test_attach_reference_guidance_never_mutates_the_trailing_user_turn(): + task = {"role": "user", "content": "ORIGINAL TASK"} + messages = [{"role": "system", "content": "sys"}, task] + moa_loop._attach_reference_guidance(messages, "REFERENCE BLOCK") + + assert messages[1] == {"role": "user", "content": "ORIGINAL TASK"} + assert messages[-1] == {"role": "user", "content": "REFERENCE BLOCK"} + assert moa_loop.peel_reference_guidance(messages, "REFERENCE BLOCK") == messages[:-1] + + +def test_prepared_aggregator_requests_share_a_byte_identical_prefix_across_iterations(monkeypatch): + calls = [] + monkeypatch.setattr(moa_loop, "call_llm", lambda **kw: calls.append(kw) or SimpleNamespace(choices=[])) + monkeypatch.setattr( + moa_loop, "_slot_runtime", + lambda slot: {"provider": "nous", "model": "openai/gpt-6-astra", "api_mode": "chat_completions"}, + ) + facade = moa_loop.MoAChatCompletions.__new__(moa_loop.MoAChatCompletions) + facade._pending_trace = None + facade._agent = None + aggregator = {"provider": "nous", "model": "openai/gpt-6-astra"} + guidance = "[Mixture of Agents reference context]\nadvice" + history = [{"role": "system", "content": "sys"}, {"role": "user", "content": "task"}] + + # Iteration 1 ends on the user task; iteration 2 replays it plus the tool round. + for messages in ( + history, + [*history, {"role": "assistant", "content": "", "tool_calls": [{"id": "1"}]}, + {"role": "tool", "tool_call_id": "1", "content": "result"}], + ): + prepared = facade.rebase_prepared_request({"guidance": guidance, "aggregator": aggregator, + "aggregator_temperature": None}, messages) + facade._call_prepared_aggregator(prepared, {"tools": [{"type": "function", "function": {"name": "lookup"}}]}) + + first, second = (c["messages"] for c in calls) + assert second[: len(first) - 1] == first[:-1] + assert second[-1] == first[-1] == {"role": "user", "content": guidance} diff --git a/website/docs/user-guide/features/mixture-of-agents.md b/website/docs/user-guide/features/mixture-of-agents.md index f983918bba..5f197f2a1d 100644 --- a/website/docs/user-guide/features/mixture-of-agents.md +++ b/website/docs/user-guide/features/mixture-of-agents.md @@ -243,7 +243,7 @@ MoA is built so the **main conversation's prompt cache is never broken**. Select Both internal call types cache normally: - **Reference models** receive a trimmed, deterministic view of the conversation (system prompt and tool transcript stripped — see the loop above). Because that view is a stable function of the stable history, a reference model's prompt prefix repeats across iterations and caches normally. References are short advisory calls with no tools. -- **The aggregator** is the acting model. The reference outputs are appended to the *end* of the latest user turn as private guidance. Because that text sits at the tail — below the entire stable prefix (system prompt + prior history) — it does not invalidate any cached prefix: the aggregator gets a cache hit on everything above the injection, and only the freshly appended tail is new. That is exactly how every normal turn behaves, where each new user message is also uncached tail tokens. +- **The aggregator** is the acting model. The reference outputs are appended as their *own* trailing user message of private guidance — never merged into your message. Because that block sits at the tail — below the entire stable prefix (system prompt + your message + prior tool history) — it does not invalidate any cached prefix: every request in a tool loop is a byte-identical extension of the previous one minus its guidance block, so the aggregator gets a cache hit on everything above the injection and only the freshly appended tail is new. That is exactly how every normal turn behaves, where each new user message is also uncached tail tokens. So MoA does not sacrifice prompt caching on either call type. Its only real cost is the extra reference calls per iteration — you pay for multiple model perspectives, not for broken caches. The long-lived conversation prefix shared with the rest of Hermes is fully intact.