fix(llm): pass LLM_EXTRA_HEADERS through ModelSettings so they reach the agent loop (#937)

tests/test_compaction.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any
 
 import pytest
 from litellm.exceptions import BadRequestError, ContextWindowExceededError, RateLimitError
+from openai.types.responses import ResponseOutputMessage, ResponseOutputText
 
 from strix.config import ContextSettings
 from strix.llm import compaction
@@ -146,17 +147,35 @@ def _patch_budget(monkeypatch: pytest.MonkeyPatch, *, keep_tokens: int, window:
     context.auto_compact = True
     settings = SimpleNamespace(
         context=context,
-        llm=SimpleNamespace(api_key=None, api_base=None, timeout=1),
+        llm=SimpleNamespace(api_key=None, api_base=None, timeout=1, extra_headers=None),
     )
     monkeypatch.setattr(compaction, "load_settings", lambda: settings)
 
 
-def _patch_summary(monkeypatch: pytest.MonkeyPatch, text: str) -> None:
-    async def fake_acompletion(**_kwargs: Any) -> Any:
-        message = SimpleNamespace(content=text)
-        return SimpleNamespace(choices=[SimpleNamespace(message=message)])
+def _model_response(text: str) -> Any:
+    chunk = ResponseOutputText(annotations=[], text=text, type="output_text")
+    message = ResponseOutputMessage(
+        id="msg", content=[chunk], role="assistant", status="completed", type="message"
+    )
+    return SimpleNamespace(output=[message])
 
-    monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
+
+def _patch_summary(
+    monkeypatch: pytest.MonkeyPatch, text: str, captured: dict[str, Any] | None = None
+) -> None:
+    class FakeModel:
+        async def get_response(self, **kwargs: Any) -> Any:
+            if captured is not None:
+                captured.update(kwargs)
+            return _model_response(text)
+
+    class FakeProvider:
+        def get_model(self, model_name: str | None) -> Any:
+            if captured is not None:
+                captured["model"] = model_name
+            return FakeModel()
+
+    monkeypatch.setattr(compaction, "StrixProvider", FakeProvider)
 
 
 @pytest.mark.asyncio
@@ -189,19 +208,38 @@ async def test_maybe_compact_rewrites_and_keeps_pairs(monkeypatch: pytest.Monkey
 async def test_maybe_compact_updates_previous_summary(monkeypatch: pytest.MonkeyPatch) -> None:
     # Window large enough to leave real room for the summary instructions.
     _patch_budget(monkeypatch, keep_tokens=30, window=4_000)
-    captured: dict[str, str] = {}
-
-    async def fake_acompletion(**kwargs: Any) -> Any:
-        captured["prompt"] = kwargs["messages"][0]["content"]
-        return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="NEW"))])
-
-    monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
+    captured: dict[str, Any] = {}
+    _patch_summary(monkeypatch, "NEW", captured)
 
     prior = compaction._checkpoint_item("OLD SUMMARY TEXT")
     session = FakeSession([prior, *_turns(12)])
 
     assert await compaction.maybe_compact(session, model="m", force=True) is True
-    assert "OLD SUMMARY TEXT" in captured["prompt"]
+    assert "OLD SUMMARY TEXT" in captured["input"]
+
+
+@pytest.mark.asyncio
+async def test_summarize_routes_through_provider_with_settings(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _patch_budget(monkeypatch, keep_tokens=30, window=4_000)
+    monkeypatch.setattr(
+        compaction,
+        "load_settings",
+        lambda: SimpleNamespace(
+            llm=SimpleNamespace(
+                api_key=None, api_base=None, timeout=1, extra_headers={"X-Feature-Key": "svc"}
+            )
+        ),
+    )
+    captured: dict[str, Any] = {}
+    _patch_summary(monkeypatch, "S", captured)
+
+    assert await compaction._summarize("litellm/openai/some-model", "p", 64) == "S"
+    assert captured["model"] == "litellm/openai/some-model"
+    settings = captured["model_settings"]
+    assert settings.extra_headers == {"X-Feature-Key": "svc"}
+    assert settings.max_tokens == 64
 
 
 def test_fit_to_tokens_truncates_oversized_text(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -233,19 +271,14 @@ def test_summary_output_tokens_capped_at_model_limit(monkeypatch: pytest.MonkeyP
 async def test_maybe_compact_bounds_summary_prompt(monkeypatch: pytest.MonkeyPatch) -> None:
     # A tiny window with a huge head must not send an oversized summary request.
     _patch_budget(monkeypatch, keep_tokens=30, window=4_000)
-    captured: dict[str, str] = {}
-
-    async def fake_acompletion(**kwargs: Any) -> Any:
-        captured["prompt"] = kwargs["messages"][0]["content"]
-        return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="S"))])
-
-    monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
+    captured: dict[str, Any] = {}
+    _patch_summary(monkeypatch, "S", captured)
     big_turns = [{"role": "user", "content": "y" * 2_000} for _ in range(50)]
     session = FakeSession(big_turns)
 
     assert await compaction.maybe_compact(session, model="m") is True
     # count_tokens==len(chars); prompt must fit the model window.
-    assert len(captured["prompt"]) <= 4_000
+    assert len(captured["input"]) <= 4_000
 
 
 @pytest.mark.asyncio
@@ -256,27 +289,27 @@ async def test_summary_request_fits_when_room_is_below_old_floor(
     instructions = len(compaction._SUMMARY_INSTRUCTIONS)
     window = instructions + 64 + 256 + 300  # summary_max(64)+slack(256)+room(300)
     _patch_budget(monkeypatch, keep_tokens=30, window=window)
-    captured: dict[str, str] = {}
-
-    async def fake_acompletion(**kwargs: Any) -> Any:
-        captured["prompt"] = kwargs["messages"][0]["content"]
-        return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="S"))])
-
-    monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
+    captured: dict[str, Any] = {}
+    _patch_summary(monkeypatch, "S", captured)
     session = FakeSession([{"role": "user", "content": "y" * 5_000} for _ in range(20)])
 
     assert await compaction.maybe_compact(session, model="m") is True
-    assert len(captured["prompt"]) <= window
+    assert len(captured["input"]) <= window
 
 
 @pytest.mark.asyncio
 async def test_maybe_compact_skips_when_summary_fails(monkeypatch: pytest.MonkeyPatch) -> None:
     _patch_budget(monkeypatch, keep_tokens=30, window=4_000)
 
-    async def fake_acompletion(**_kwargs: Any) -> Any:
-        raise RuntimeError("boom")
+    class BoomModel:
+        async def get_response(self, **_kwargs: Any) -> Any:
+            raise RuntimeError("boom")
 
-    monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
+    class BoomProvider:
+        def get_model(self, _model_name: str | None) -> Any:
+            return BoomModel()
+
+    monkeypatch.setattr(compaction, "StrixProvider", BoomProvider)
     session = FakeSession(_turns(12))
     before = await session.get_items()
 
@@ -290,17 +323,11 @@ async def test_maybe_compact_skips_when_no_room_to_summarise(
 ) -> None:
     # No room for any head -> no (doomed) summary is attempted.
     _patch_budget(monkeypatch, keep_tokens=30, window=200)
-    called = False
-
-    async def fake_acompletion(**_kwargs: Any) -> Any:
-        nonlocal called
-        called = True
-        return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="S"))])
-
-    monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
+    captured: dict[str, Any] = {}
+    _patch_summary(monkeypatch, "S", captured)
     session = FakeSession(_turns(12))
     before = await session.get_items()
 
     assert await compaction.maybe_compact(session, model="m", force=True) is False
-    assert called is False
+    assert not captured
     assert await session.get_items() == before

tests/test_dedupe_model.py

@@ -44,6 +44,38 @@ def test_dedupe_endpoint_sent_per_call() -> None:
     assert (settings.extra_args or {})["api_key"] == "dedupe-key"
 
 
+def test_dedicated_dedupe_model_uses_own_headers_not_main() -> None:
+    dedupe = DedupeSettings(
+        STRIX_DEDUPE_MODEL="deepseek/cheap",
+        DEDUPE_LLM_EXTRA_HEADERS={"X-Dedupe": "yes"},
+    )
+    settings = _dedupe_model_settings(dedupe, "deepseek/cheap", 300)
+    assert settings.extra_headers == {"X-Dedupe": "yes"}
+
+
+def test_dedicated_dedupe_model_gets_no_main_headers_by_default(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("LLM_EXTRA_HEADERS", json.dumps({"X-Main": "secret"}))
+    loader._cached = None
+    try:
+        dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap")
+        settings = _dedupe_model_settings(dedupe, "deepseek/cheap", 300)
+        assert settings.extra_headers is None
+    finally:
+        loader._cached = None
+
+
+def test_fallback_dedupe_inherits_main_headers(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("LLM_EXTRA_HEADERS", json.dumps({"X-Main": "svc"}))
+    loader._cached = None
+    try:
+        settings = _dedupe_model_settings(DedupeSettings(), "openai/main-model", 300)
+        assert settings.extra_headers == {"X-Main": "svc"}
+    finally:
+        loader._cached = None
+
+
 def test_dedupe_defaults_are_empty() -> None:
     settings = DedupeSettings()
     assert settings.model is None

tests/test_inputs.py

@@ -272,6 +272,30 @@ def test_make_model_settings_omits_timeout_when_unset() -> None:
     assert settings.extra_args is None
 
 
+def test_make_model_settings_sets_extra_headers() -> None:
+    settings = make_model_settings(
+        "none",
+        model_name="openai/some-model",
+        extra_headers={"X-Feature-Key": "svc", "X-Tenant": "acme"},
+    )
+
+    assert settings.extra_headers == {"X-Feature-Key": "svc", "X-Tenant": "acme"}
+
+
+def test_make_model_settings_omits_extra_headers_when_unset() -> None:
+    assert make_model_settings("none", model_name="gpt-4o").extra_headers is None
+
+
+def test_make_model_settings_extra_headers_survive_reasoning_resolve() -> None:
+    settings = make_model_settings(
+        "high",
+        model_name="openai/o3",
+        extra_headers={"X-Feature-Key": "svc"},
+    )
+
+    assert settings.extra_headers == {"X-Feature-Key": "svc"}
+
+
 def test_make_model_settings_timeout_survives_reasoning_resolve() -> None:
     # Reasoning is resolved via ModelSettings.resolve(); the timeout in extra_args
     # must not be dropped when a reasoning override is merged in.

tests/test_runner_rate_limit.py

@@ -40,6 +40,7 @@ async def test_persistent_rate_limit_stops_gracefully(
             force_required_tool_choice=False,
             timeout=300,
             prompt_cache=True,
+            extra_headers=None,
         ),
         runtime=types.SimpleNamespace(max_context_images=3),
     )

tests/test_runner_root_prompt.py

@@ -48,6 +48,7 @@ def _patch_engine_scaffold(
             force_required_tool_choice=False,
             timeout=300,
             prompt_cache=True,
+            extra_headers=None,
         ),
         runtime=types.SimpleNamespace(max_context_images=3),
     )