fix(moa): pass Copilot initiator header to advisors

This commit is contained in:
dyreckt 2026-07-07 10:40:44 -04:00 committed by Teknium
parent 0749cac7a1
commit 4c66307c36
3 changed files with 113 additions and 0 deletions

View file

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

View file

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

View file

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