Skip to content
22 changes: 12 additions & 10 deletions mcp_server/services/llm_tiers.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,17 @@ def apply_mcp_servers_enabled_form(form_data: Dict[str, str]) -> None:
form_data[MCP_SERVERS_ENABLED_ENV] = value.strip()


def apply_llm_tier_overrides(form_data: Dict[str, str]) -> None:
"""Inject LLM tier overrides from the MCP environment into form_data (in-place).
def llm_tier_overrides_from_env() -> Dict[str, str]:
"""Return the LLM tier overrides present in the process environment.

Reads LLM_HIGH, LLM_MEDIUM, LLM_LOW from the process environment (set via
mcp.json) and adds any that are present to the form_data dict that will be
sent to the backend API. Keys not set in the environment are left out so
the backend falls back to its own defaults.
Reads LLM_HIGH, LLM_MEDIUM, LLM_LOW (set via mcp.json) and returns only the
ones actually set, so absent tiers let the backend fall back to its own
defaults. Pure — no mutation — so callers can pass the result directly.
"""
for key in LLM_TIER_KEYS:
value = os.environ.get(key)
if value:
form_data[key] = value
return {key: value for key in LLM_TIER_KEYS if (value := os.environ.get(key))}


def apply_llm_tier_overrides(form_data: Dict[str, str]) -> None:
"""In-place variant of :func:`llm_tier_overrides_from_env` for callers that
accumulate several groups of fields into one ``form_data`` dict."""
form_data.update(llm_tier_overrides_from_env())
37 changes: 30 additions & 7 deletions mcp_server/services/validate_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@

from fastmcp import Context

from services.llm_tiers import apply_llm_tier_overrides
from services.llm_tiers import LLM_TIER_KEYS, llm_tier_overrides_from_env
from services.run_generation_precheck import RejectionCode
from services.specflow_backend import post_form_data_to_backend

Expand All @@ -27,20 +27,43 @@
VALIDATE_MODELS_ENDPOINT = "/api/v1/models/validate"


async def request_model_validation() -> dict[str, Any]:
"""Call the backend validate-models endpoint with the configured LLM_* overrides.
async def _post_validation(form_data: dict[str, str]) -> dict[str, Any]:
"""POST ``form_data`` to the validate-models endpoint and parse the response.

Returns the parsed ValidateModelsResponse dict. Raises on transport/parse errors
so callers can stay permissive (connect/run gate) on infrastructure failures.
The single call site for the endpoint so URL/timeout/parsing live in one place.
"""
form_data: dict[str, str] = {}
apply_llm_tier_overrides(form_data) # injects LLM_HIGH/MEDIUM/LOW from env when set
raw = await post_form_data_to_backend(
VALIDATE_MODELS_ENDPOINT, form_data, timeout_seconds=30.0
)
return json.loads(raw)


async def request_model_validation() -> dict[str, Any]:
"""Call the backend validate-models endpoint with the configured LLM_* overrides.

Reads LLM_HIGH/MEDIUM/LOW from the MCP *process env*. Returns the parsed
ValidateModelsResponse dict. Raises on transport/parse errors so callers can
stay permissive (connect/run gate) on infrastructure failures.
"""
return await _post_validation(llm_tier_overrides_from_env())


async def request_model_validation_for(tier_values: dict[str, str]) -> dict[str, Any]:
"""Validate explicit LLM_* ``tier_values`` (e.g. from the local config block).

Unlike ``request_model_validation`` (which reads the process env), this takes
the values directly — the TUI edits ``mcp-config.json`` rather than its own
environment, so it must forward what the file holds. Blank tiers are omitted
so the backend falls back to its own default for that tier.
"""
form_data = {
key: value.strip()
for key in LLM_TIER_KEYS
if (value := (tier_values.get(key) or "").strip())
}
return await _post_validation(form_data)


def invalid_models(response: dict[str, Any]) -> list[dict[str, Any]]:
"""Flatten the response into the list of confidently-invalid models (with tier)."""
out: list[dict[str, Any]] = []
Expand Down
26 changes: 26 additions & 0 deletions mcp_server/tests/test_llm_tiers.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,35 @@
MCP_SERVERS_ENABLED_ENV,
apply_llm_tier_overrides,
apply_mcp_servers_enabled_form,
llm_tier_overrides_from_env,
)


class TestLlmTierOverridesFromEnv:
def test_returns_only_keys_present_in_env(self, monkeypatch):
monkeypatch.setenv("LLM_HIGH", "claude-opus")
monkeypatch.setenv("LLM_LOW", "claude-haiku")
monkeypatch.delenv("LLM_MEDIUM", raising=False)
assert llm_tier_overrides_from_env() == {
"LLM_HIGH": "claude-opus",
"LLM_LOW": "claude-haiku",
}

def test_empty_when_no_env(self, monkeypatch):
for key in LLM_TIER_KEYS:
monkeypatch.delenv(key, raising=False)
assert llm_tier_overrides_from_env() == {}

def test_apply_is_the_in_place_form_of_the_pure_helper(self, monkeypatch):
monkeypatch.setenv("LLM_HIGH", "claude-opus")
monkeypatch.delenv("LLM_MEDIUM", raising=False)
monkeypatch.delenv("LLM_LOW", raising=False)
form: dict[str, str] = {"spec_path": "specs"}
apply_llm_tier_overrides(form)
# Wrapper merges the pure result, preserving unrelated keys.
assert form == {"spec_path": "specs", "LLM_HIGH": "claude-opus"}


class TestApplyLlmTierOverrides:
def test_injects_only_keys_present_in_env(self, monkeypatch):
monkeypatch.setenv("LLM_HIGH", "claude-opus")
Expand Down
Loading