From 32a4faa2d5b0eb66a6c85bf9be6436c5afc76318 Mon Sep 17 00:00:00 2001 From: Janig88 Date: Sat, 4 Jul 2026 20:28:25 +0300 Subject: [PATCH] fix(auxiliary): honor max_tokens for MoA reference/aggregator tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- agent/auxiliary_client.py | 9 ++- contributors/emails/janig88@gmail.com | 2 + tests/agent/test_auxiliary_client.py | 95 +++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 contributors/emails/janig88@gmail.com diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index c678bc2a3796..832559aa6381 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -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): diff --git a/contributors/emails/janig88@gmail.com b/contributors/emails/janig88@gmail.com new file mode 100644 index 000000000000..2eb927a7e08c --- /dev/null +++ b/contributors/emails/janig88@gmail.com @@ -0,0 +1,2 @@ +Janig88 +# PR #58402 salvage diff --git a/tests/agent/test_auxiliary_client.py b/tests/agent/test_auxiliary_client.py index 9f61ff9668db..fc4f8f562dc2 100644 --- a/tests/agent/test_auxiliary_client.py +++ b/tests/agent/test_auxiliary_client.py @@ -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):