fix(moa): keep the aggregator prefix byte-stable across tool-loop iterations

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.
This commit is contained in:
teknium1
2026-09-16 10:18:06 -07:00
committed by Teknium
parent 095de7ffd0
commit f3e62619c7
4 changed files with 68 additions and 31 deletions

View File

@@ -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

View File

@@ -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)

View File

@@ -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}

View File

@@ -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.