fix(auxiliary): honor max_tokens for MoA reference/aggregator tasks

PR #56756 added reference_max_tokens to cap MoA advisor output and cut
turn latency. The value is correctly threaded through five layers of MoA
code (moa_config → conversation_loop → aggregate_moa_context →
_run_references_parallel → _run_reference → call_llm(task='moa_reference',
max_tokens=800, ...)).

However, _build_call_kwargs() in auxiliary_client.py silently drops
max_tokens for all OpenAI-compatible providers (PR #34845, which fixed
endpoints and NVIDIA NIM keep it. This means reference_max_tokens never
reached the API for the vast majority of providers.

The bug affects every OpenAI-compatible MoA reference/aggregator slot:
Z.AI (coding plan), OpenRouter, OpenAI, GitHub Copilot, and local
providers. Only Anthropic-compat endpoints (MiniMax, /anthropic URLs)
worked — by coincidence, not MoA-aware design.

Fix: thread the 'task' parameter through all six _build_call_kwargs()
call sites. When task starts with 'moa_', max_tokens is always included
in the request kwargs regardless of provider. Non-MoA auxiliary tasks
(compression, titles, vision, etc.) keep PR #34845 behavior unchanged.

Verified end-to-end:
- Z.AI GLM-5.2 with max_tokens=50 → returned exactly 50 tokens
- Z.AI GLM-5.2 with max_tokens=20 → returned exactly 20 tokens
- Z.AI GLM-5.2 uncapped → returned 315 tokens
- 7 new regression tests covering 4 providers, Anthropic wire, non-MoA
  tasks, and prefix-matching boundary
- 288 auxiliary_client tests pass (was 281, +7 new), 84 MoA tests pass
- Zero regressions
This commit is contained in:
Janig88 2026-07-04 20:28:25 +03:00 committed by Teknium
parent cc1725cbe5
commit 32a4faa2d5
3 changed files with 104 additions and 2 deletions

View file

@ -3683,6 +3683,7 @@ def _retry_same_provider_sync(
extra_body=effective_extra_body,
reasoning_config=reasoning_config,
base_url=retry_base or resolved_base_url,
task=task,
)
if _is_anthropic_compat_endpoint(resolved_provider, retry_base):
retry_kwargs["messages"] = _convert_openai_images_to_anthropic(retry_kwargs["messages"])
@ -3742,6 +3743,7 @@ async def _retry_same_provider_async(
extra_body=effective_extra_body,
reasoning_config=reasoning_config,
base_url=retry_base or resolved_base_url,
task=task,
)
if _is_anthropic_compat_endpoint(resolved_provider, retry_base):
retry_kwargs["messages"] = _convert_openai_images_to_anthropic(retry_kwargs["messages"])
@ -6776,6 +6778,7 @@ def _build_call_kwargs(
extra_body: Optional[dict] = None,
reasoning_config: Optional[dict] = None,
base_url: Optional[str] = None,
task: Optional[str] = None,
) -> dict:
"""Build kwargs for .chat.completions.create() with model/provider adjustments."""
kwargs: Dict[str, Any] = {
@ -6830,9 +6833,11 @@ def _build_call_kwargs(
_provider_norm in {"nvidia", "nvidia-nim", "nim", "build-nvidia", "nemotron"}
or base_url_host_matches(_effective_base, "integrate.api.nvidia.com")
)
_is_moa = bool(task) and str(task).startswith("moa_")
if (
_is_anthropic_compat_endpoint(provider, _effective_base)
or _is_nvidia_nim
or _is_moa
):
kwargs["max_tokens"] = max_tokens
@ -7203,7 +7208,7 @@ def call_llm(
temperature=temperature, max_tokens=max_tokens,
tools=tools, timeout=effective_timeout, extra_body=effective_extra_body,
reasoning_config=reasoning_config,
base_url=_base_info or resolved_base_url)
base_url=_base_info or resolved_base_url, task=task)
# Convert image blocks for Anthropic-compatible endpoints (e.g. MiniMax)
_client_base = str(getattr(client, "base_url", "") or "")
@ -7819,7 +7824,7 @@ async def async_call_llm(
temperature=temperature, max_tokens=max_tokens,
tools=tools, timeout=effective_timeout, extra_body=effective_extra_body,
reasoning_config=reasoning_config,
base_url=_client_base or resolved_base_url)
base_url=_client_base or resolved_base_url, task=task)
# Convert image blocks for Anthropic-compatible endpoints (e.g. MiniMax)
if _is_anthropic_compat_endpoint(resolved_provider, _client_base):

View file

@ -0,0 +1,2 @@
Janig88
# PR #58402 salvage

View file

@ -656,6 +656,101 @@ class TestBuildCallKwargsMaxTokens:
)
assert kwargs["max_tokens"] == 4096
# ── MoA task should honor max_tokens on ALL providers (#reference_max_tokens) ──
@pytest.mark.parametrize(
"provider,model,base_url",
[
("zai", "glm-5.2", "https://api.z.ai/api/coding/paas/v4"),
("openrouter", "deepseek/deepseek-v4-flash:nitro", "https://openrouter.ai/api/v1"),
("copilot", "gpt-5.5", "https://api.githubcopilot.com"),
("nous", "hermes-4", "https://inference-api.nousresearch.com/v1"),
],
)
def test_moa_task_sends_max_tokens_on_openai_compatible(self, provider, model, base_url):
"""MoA reference/aggregator tasks must honor max_tokens regardless of provider.
The ``reference_max_tokens`` config option (PR #56756) caps advisor output
to reduce turn latency. Before the fix, ``_build_call_kwargs`` silently
dropped the value for OpenAI-compatible providers (PR #34845), so the cap
never reached the API. With the ``task`` parameter threaded through,
any task starting with ``moa_`` must include ``max_tokens`` in kwargs.
"""
from agent.auxiliary_client import _build_call_kwargs
kwargs = _build_call_kwargs(
provider=provider,
model=model,
messages=[{"role": "user", "content": "hi"}],
max_tokens=800,
base_url=base_url,
task="moa_reference",
)
assert kwargs["max_tokens"] == 800
def test_moa_task_sends_max_tokens_on_anthropic_wire(self):
"""MoA tasks on Anthropic-compat endpoints keep max_tokens (unchanged behavior)."""
from agent.auxiliary_client import _build_call_kwargs
kwargs = _build_call_kwargs(
provider="minimax",
model="minimax-m2",
messages=[{"role": "user", "content": "hi"}],
max_tokens=600,
base_url="https://api.minimax.io/v1",
task="moa_aggregator",
)
assert kwargs["max_tokens"] == 600
def test_non_moa_tasks_still_omit_max_tokens(self):
"""Regression guard: compression/titles/vision keep PR #34845 behavior."""
from agent.auxiliary_client import _build_call_kwargs
for task in ("compression", "vision", "title_generation", None, ""):
kwargs = _build_call_kwargs(
provider="openrouter",
model="deepseek/deepseek-v4-flash:nitro",
messages=[{"role": "user", "content": "hi"}],
max_tokens=800,
base_url="https://openrouter.ai/api/v1",
task=task,
)
assert "max_tokens" not in kwargs, f"max_tokens should be dropped for task={task!r}"
def test_moa_prefix_matching(self):
"""Only tasks prefixed with 'moa_' trigger the cap — not arbitrary task names."""
from agent.auxiliary_client import _build_call_kwargs
# 'moa_reference' → honored
kw = _build_call_kwargs(
provider="zai", model="glm-5.2",
messages=[{"role": "user", "content": "hi"}],
max_tokens=500,
base_url="https://api.z.ai/api/coding/paas/v4",
task="moa_reference",
)
assert kw["max_tokens"] == 500
# 'moa_xyz' → honored (prefix match)
kw2 = _build_call_kwargs(
provider="zai", model="glm-5.2",
messages=[{"role": "user", "content": "hi"}],
max_tokens=500,
base_url="https://api.z.ai/api/coding/paas/v4",
task="moa_custom_future",
)
assert kw2["max_tokens"] == 500
# 'mopha_reference' (similar but not moa_) → dropped
kw3 = _build_call_kwargs(
provider="zai", model="glm-5.2",
messages=[{"role": "user", "content": "hi"}],
max_tokens=500,
base_url="https://api.z.ai/api/coding/paas/v4",
task="mopha_reference",
)
assert "max_tokens" not in kw3
class TestNousTagsScoping:
def test_tags_injected_when_provider_is_nous(self, monkeypatch):