From d4e58b2cd0753870276677cfc182905f04ad6570 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 19:38:06 -0700 Subject: [PATCH] fix(llm): pass LLM_EXTRA_HEADERS through ModelSettings so they reach the agent loop (#937) --- docs/advanced/configuration.mdx | 6 ++ strix/config/settings.py | 4 ++ strix/core/inputs.py | 2 + strix/core/runner.py | 1 + strix/interface/main.py | 24 ++++++- strix/interface/tui/app.py | 3 +- strix/llm/compaction.py | 56 ++++++++++++---- strix/report/dedupe.py | 11 ++- tests/test_compaction.py | 111 +++++++++++++++++++------------ tests/test_dedupe_model.py | 32 +++++++++ tests/test_inputs.py | 24 +++++++ tests/test_runner_rate_limit.py | 1 + tests/test_runner_root_prompt.py | 1 + 13 files changed, 215 insertions(+), 61 deletions(-) diff --git a/docs/advanced/configuration.mdx b/docs/advanced/configuration.mdx index b446fb67..af98b8b8 100644 --- a/docs/advanced/configuration.mdx +++ b/docs/advanced/configuration.mdx @@ -63,6 +63,12 @@ affecting the agents that do the actual testing. model runs on a different endpoint than the main model. + + Optional JSON object of extra HTTP headers sent on every deduplication-model + request, e.g. `{"X-Feature-Key":"value"}`. A dedicated dedupe model never + inherits `LLM_EXTRA_HEADERS`; set this when its endpoint needs custom headers. + + Reasoning effort for the deduplication model. Defaults to the model's own baseline when unset. diff --git a/strix/config/settings.py b/strix/config/settings.py index 31f479b0..016a8ad9 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -61,6 +61,10 @@ class DedupeSettings(BaseSettings): ) api_key: str | None = Field(default=None, alias="DEDUPE_LLM_API_KEY") api_base: str | None = Field(default=None, alias="DEDUPE_LLM_API_BASE") + extra_headers: dict[str, str] | None = Field( + default=None, + alias="DEDUPE_LLM_EXTRA_HEADERS", + ) class ContextSettings(BaseSettings): diff --git a/strix/core/inputs.py b/strix/core/inputs.py index aef1fe13..34a2d4b3 100644 --- a/strix/core/inputs.py +++ b/strix/core/inputs.py @@ -132,12 +132,14 @@ def make_model_settings( force_required_tool_choice: bool = False, request_timeout: float | None = None, prompt_cache: bool = True, + extra_headers: dict[str, str] | None = None, ) -> ModelSettings: model_settings = ModelSettings( parallel_tool_calls=False, retry=DEFAULT_MODEL_RETRY, include_usage=True, extra_args=request_timeout_extra_args(request_timeout), + extra_headers=dict(extra_headers) if extra_headers else None, ) if ( reasoning_effort is not None diff --git a/strix/core/runner.py b/strix/core/runner.py index c5f51b15..01725cab 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -250,6 +250,7 @@ async def run_strix_scan( force_required_tool_choice=settings.llm.force_required_tool_choice, request_timeout=settings.llm.timeout, prompt_cache=settings.llm.prompt_cache, + extra_headers=settings.llm.extra_headers, ) run_config = RunConfig( model=resolved_model, diff --git a/strix/interface/main.py b/strix/interface/main.py index 0fbcbf48..4d88beda 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -31,7 +31,7 @@ from strix.config.models import ( is_known_openai_bare_model, is_recommended_or_frontier_model, ) -from strix.core.inputs import DEFAULT_MAX_TURNS +from strix.core.inputs import DEFAULT_MAX_TURNS, make_model_settings from strix.core.paths import run_dir_for, runtime_state_dir from strix.interface.cli import run_cli from strix.interface.tui import run_tui @@ -382,7 +382,13 @@ async def warm_up_llm(show_model_warning: bool = True) -> None: model.get_response( system_instructions="You are a helpful assistant.", input="Reply with just 'OK'.", - model_settings=ModelSettings(), + model_settings=make_model_settings( + None, + model_name=raw_model, + request_timeout=llm.timeout, + prompt_cache=False, + extra_headers=llm.extra_headers, + ), tools=[], output_schema=None, handoffs=[], @@ -404,7 +410,19 @@ async def warm_up_llm(show_model_warning: bool = True) -> None: # Match the runtime path: send the dedupe key/endpoint per call so a # separate-provider dedupe model authenticates during warm-up too. deduper_extra = _dedupe_extra_args(settings.dedupe) - deduper_settings = ModelSettings(extra_args=deduper_extra or None) + # A dedicated dedupe model may route to another provider, which must + # never receive the main endpoint's headers; it has its own + # DEDUPE_LLM_EXTRA_HEADERS. + deduper_settings = make_model_settings( + None, + model_name=dedupe_model, + request_timeout=llm.timeout, + prompt_cache=False, + extra_headers=settings.dedupe.extra_headers, + ) + if deduper_extra: + merged = {**(deduper_settings.extra_args or {}), **deduper_extra} + deduper_settings = deduper_settings.resolve(ModelSettings(extra_args=merged)) await asyncio.wait_for( deduper.get_response( system_instructions="You are a helpful assistant.", diff --git a/strix/interface/tui/app.py b/strix/interface/tui/app.py index cc5a36c4..14bc6cb1 100644 --- a/strix/interface/tui/app.py +++ b/strix/interface/tui/app.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, ClassVar if TYPE_CHECKING: + from pygments.token import _TokenType from textual.timer import Timer from rich.align import Align @@ -352,7 +353,7 @@ class VulnerabilityDetailScreen(ModalScreen): # type: ignore[misc] if not token_value: continue color = None - tt = token_type + tt: _TokenType | None = token_type while tt: if tt in colors: color = colors[tt] diff --git a/strix/llm/compaction.py b/strix/llm/compaction.py index da612ff2..09d36454 100644 --- a/strix/llm/compaction.py +++ b/strix/llm/compaction.py @@ -12,15 +12,20 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any -import litellm +from agents.model_settings import ModelSettings +from agents.models.interface import ModelTracing from litellm.exceptions import BadRequestError, ContextWindowExceededError +from openai.types.responses import ResponseOutputMessage, ResponseOutputText from strix.config import load_settings +from strix.config.models import StrixProvider +from strix.core.inputs import make_model_settings from strix.core.sessions import replace_session_items, session_write_lock from strix.llm.context_budget import context_window, count_tokens, output_limit if TYPE_CHECKING: + from agents.items import ModelResponse from agents.memory import Session @@ -268,26 +273,53 @@ def _checkpoint_item(summary: str) -> dict[str, Any]: } +def _extract_text(response: ModelResponse) -> str: + parts: list[str] = [] + for item in response.output: + if not isinstance(item, ResponseOutputMessage): + continue + parts.extend( + chunk.text + for chunk in item.content + if isinstance(chunk, ResponseOutputText) and chunk.text + ) + return "".join(parts) + + async def _summarize(model: str, prompt: str, max_tokens: int) -> str | None: llm = load_settings().llm + model_settings = make_model_settings( + None, + model_name=model, + request_timeout=llm.timeout, + prompt_cache=False, + extra_headers=llm.extra_headers, + ).resolve(ModelSettings(max_tokens=max_tokens)) try: - response = await litellm.acompletion( - model=model, - messages=[{"role": "user", "content": prompt}], - max_tokens=max_tokens, - api_key=llm.api_key, - api_base=llm.api_base, - timeout=llm.timeout, + response = ( + await StrixProvider() + .get_model(model) + .get_response( + system_instructions=None, + input=prompt, + model_settings=model_settings, + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + previous_response_id=None, + conversation_id=None, + prompt=None, + ) ) except Exception: logger.exception("compaction summary call failed for model %s", model) return None - try: - content = response.choices[0].message.content - except (AttributeError, IndexError, KeyError): + content = _extract_text(response).strip() + if not content: logger.warning("compaction summary returned no content") return None - return content.strip() if isinstance(content, str) and content.strip() else None + return content async def maybe_compact( diff --git a/strix/report/dedupe.py b/strix/report/dedupe.py index b47ed5a4..23db066f 100644 --- a/strix/report/dedupe.py +++ b/strix/report/dedupe.py @@ -51,17 +51,24 @@ def _dedupe_extra_args(dedupe: DedupeSettings) -> dict[str, str]: def _dedupe_model_settings( dedupe: DedupeSettings, model_name: str, request_timeout: float | None ) -> ModelSettings: + llm = load_settings().llm settings = make_model_settings( dedupe.reasoning_effort, model_name=model_name, force_required_tool_choice=False, request_timeout=request_timeout, + # The main model's headers apply only when dedupe falls back to the main + # model; a dedicated dedupe model may route to another provider, which + # must never receive the main endpoint's credentials. A dedicated model + # gets its own DEDUPE_LLM_EXTRA_HEADERS instead. + extra_headers=dedupe.extra_headers if dedupe.model else llm.extra_headers, ) extra = _dedupe_extra_args(dedupe) if extra: settings = settings.resolve(ModelSettings(extra_args=extra)) return settings + DEDUPE_SYSTEM_PROMPT = """You are an expert vulnerability report deduplication judge. Your task is to determine if a candidate vulnerability report describes the SAME vulnerability as any existing report. @@ -347,9 +354,7 @@ async def check_duplicate( response = await model.get_response( system_instructions=DEDUPE_SYSTEM_PROMPT, input=user_msg, - model_settings=_dedupe_model_settings( - dedupe, resolved_model, settings.llm.timeout - ), + model_settings=_dedupe_model_settings(dedupe, resolved_model, settings.llm.timeout), tools=[], output_schema=None, handoffs=[], diff --git a/tests/test_compaction.py b/tests/test_compaction.py index da9a2361..2d73d9ac 100644 --- a/tests/test_compaction.py +++ b/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 diff --git a/tests/test_dedupe_model.py b/tests/test_dedupe_model.py index 0a0f7654..b17946e9 100644 --- a/tests/test_dedupe_model.py +++ b/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 diff --git a/tests/test_inputs.py b/tests/test_inputs.py index da914879..871ea149 100644 --- a/tests/test_inputs.py +++ b/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. diff --git a/tests/test_runner_rate_limit.py b/tests/test_runner_rate_limit.py index 482c4730..061ad3c5 100644 --- a/tests/test_runner_rate_limit.py +++ b/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), ) diff --git a/tests/test_runner_root_prompt.py b/tests/test_runner_root_prompt.py index 56d7caa6..cd4d4ac8 100644 --- a/tests/test_runner_root_prompt.py +++ b/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), )