diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index f475369d1a10..dc5c2648f7a0 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -3648,6 +3648,7 @@ def _retry_same_provider_sync( effective_timeout: float, effective_extra_body: dict, reasoning_config: Optional[dict], + extra_headers: Optional[Dict[str, str]] = None, ) -> Any: if task == "vision": _, retry_client, retry_model = resolve_vision_provider_client( @@ -3685,6 +3686,11 @@ def _retry_same_provider_sync( base_url=retry_base or resolved_base_url, task=task, ) + # Preserve per-request attribution headers (e.g. Copilot's + # ``x-initiator: user``) across the rebuilt-client retry — dropping them + # here would let a recovery retry silently lose capability gating (#60293). + if extra_headers: + retry_kwargs["extra_headers"] = dict(extra_headers) if _is_anthropic_compat_endpoint(resolved_provider, retry_base): retry_kwargs["messages"] = _convert_openai_images_to_anthropic(retry_kwargs["messages"]) return _validate_llm_response( @@ -3708,6 +3714,7 @@ async def _retry_same_provider_async( effective_timeout: float, effective_extra_body: dict, reasoning_config: Optional[dict], + extra_headers: Optional[Dict[str, str]] = None, ) -> Any: if task == "vision": _, retry_client, retry_model = resolve_vision_provider_client( @@ -3745,6 +3752,10 @@ async def _retry_same_provider_async( base_url=retry_base or resolved_base_url, task=task, ) + # Preserve per-request attribution headers across the rebuilt-client + # retry — see the sync variant above (#60293). + if extra_headers: + retry_kwargs["extra_headers"] = dict(extra_headers) if _is_anthropic_compat_endpoint(resolved_provider, retry_base): retry_kwargs["messages"] = _convert_openai_images_to_anthropic(retry_kwargs["messages"]) return _validate_llm_response( @@ -7108,6 +7119,7 @@ def call_llm( timeout: float = None, extra_body: dict = None, reasoning_config: Optional[dict] = None, + extra_headers: Optional[Dict[str, str]] = None, api_mode: str = None, stream: bool = False, stream_options: dict = None, @@ -7133,6 +7145,9 @@ def call_llm( extra_body: Additional request body fields. reasoning_config: Optional Hermes reasoning config for direct model calls such as MoA reference/aggregator slots. + extra_headers: Additional per-request HTTP headers. These override + client-level defaults for providers that gate capabilities on + request attribution (for example Copilot's ``x-initiator``). stream: When True, return the raw SDK streaming iterator instead of a validated complete response. The caller is responsible for consuming chunks (and for any fallback). Used by the MoA aggregator so its @@ -7246,6 +7261,8 @@ def call_llm( tools=tools, timeout=effective_timeout, extra_body=effective_extra_body, reasoning_config=reasoning_config, base_url=_base_info or resolved_base_url, task=task) + if extra_headers: + kwargs["extra_headers"] = dict(extra_headers) # Convert image blocks for Anthropic-compatible endpoints (e.g. MiniMax) _client_base = str(getattr(client, "base_url", "") or "") @@ -7499,6 +7516,7 @@ def call_llm( effective_timeout=effective_timeout, effective_extra_body=effective_extra_body, reasoning_config=reasoning_config, + extra_headers=extra_headers, ) # ── Same-provider credential-pool recovery ───────────────────── @@ -7542,6 +7560,7 @@ def call_llm( effective_timeout=effective_timeout, effective_extra_body=effective_extra_body, reasoning_config=reasoning_config, + extra_headers=extra_headers, ) except Exception as retry2_err: # The rotated key also hit a quota/auth wall. Mark it diff --git a/agent/moa_loop.py b/agent/moa_loop.py index b742c743ce05..8522750cd04b 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -338,12 +338,32 @@ def _run_reference( # reference model have its own output cap independently. _slot_max_tokens: int | None = slot.get("max_tokens") _effective_max_tokens = _slot_max_tokens if _slot_max_tokens is not None else max_tokens + extra_headers = None + # Normalize provider aliases (github, github-copilot, github-models, + # ...) through the auxiliary client's canonical alias table so slot + # configs that spell Copilot differently still get the header. + from agent.auxiliary_client import _normalize_aux_provider + + if _normalize_aux_provider(str(runtime.get("provider") or "")) in ( + "copilot", + "copilot-acp", + ): + # Copilot Pro/Pro+ gates some premium chat models on request + # attribution. The main agent marks the first API request of a + # user turn as ``x-initiator: user``; MoA reference fan-out is also + # directly serving the user's current turn, not a background agent + # task, so mirror that header here. Without it, Claude/Gemini + # Copilot advisors can be rejected as unavailable to the + # ``copilot-language-server`` integrator even though standalone + # Copilot calls work. + extra_headers = {"x-initiator": "user"} response = call_llm( task="moa_reference", messages=messages, temperature=temperature, max_tokens=_effective_max_tokens, reasoning_config=_slot_reasoning_config(slot), + extra_headers=extra_headers, **runtime, ) usage = CanonicalUsage() diff --git a/tests/run_agent/test_moa_loop_mode.py b/tests/run_agent/test_moa_loop_mode.py index c41dca2511db..2ff06f79b296 100644 --- a/tests/run_agent/test_moa_loop_mode.py +++ b/tests/run_agent/test_moa_loop_mode.py @@ -395,6 +395,80 @@ def test_moa_provider_backed_slot_survives_aux_resolution(monkeypatch, provider) assert api_key == f"token-for-{provider}" +def test_moa_copilot_reference_forwards_user_initiator_header(monkeypatch): + """Copilot MoA advisors must carry the same user-turn attribution as main calls. + + Copilot Pro/Pro+ gates some premium chat models on the ``x-initiator`` + request header. MoA references are direct fan-out for the user's current + turn, so Copilot advisors need ``x-initiator: user`` rather than inheriting + the Copilot language-server default attribution. + """ + from agent import moa_loop + + calls = [] + + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda _slot: { + "provider": "copilot", + "model": "claude-sonnet-4.6", + "api_mode": "chat_completions", + "base_url": "https://api.githubcopilot.com", + "api_key": "copilot-token", + }, + ) + + def fake_call_llm(**kwargs): + calls.append(kwargs) + return _response("copilot advice") + + monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm) + + _label, text, _acct = moa_loop._run_reference( + {"provider": "copilot", "model": "claude-sonnet-4.6"}, + [{"role": "user", "content": "solve this"}], + ) + + assert text == "copilot advice" + assert calls[0]["task"] == "moa_reference" + assert calls[0]["extra_headers"] == {"x-initiator": "user"} + + +def test_moa_non_copilot_reference_does_not_forward_initiator_header(monkeypatch): + """The Copilot attribution header must stay scoped to Copilot advisors.""" + from agent import moa_loop + + calls = [] + + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda _slot: { + "provider": "openrouter", + "model": "anthropic/claude-sonnet-4.6", + "api_mode": "chat_completions", + "base_url": "https://openrouter.ai/api/v1", + "api_key": "openrouter-token", + }, + ) + + def fake_call_llm(**kwargs): + calls.append(kwargs) + return _response("openrouter advice") + + monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm) + + _label, text, _acct = moa_loop._run_reference( + {"provider": "openrouter", "model": "anthropic/claude-sonnet-4.6"}, + [{"role": "user", "content": "solve this"}], + ) + + assert text == "openrouter advice" + assert calls[0]["task"] == "moa_reference" + assert calls[0]["extra_headers"] is None + + def test_moa_slot_runtime_falls_back_on_resolution_error(monkeypatch): """A slot whose provider can't be resolved still attempts the call with the bare provider/model rather than aborting the whole MoA turn."""