fix(proxy): strip tool config from health check probes (#31008)

Health checks reuse the deployment's litellm_params, so a deployment with extra_body.tools (e.g. openrouter:web_search) sent those server tools on the trivial probe and OpenRouter 500'd it, while real completions on the same deployment succeeded. _update_litellm_params_for_health_check now works on a copy and strips tools and tool_choice from both the top level and extra_body before probing, since a trivial probe never needs tool configuration
This commit is contained in:
AyushGupta-Code 2026-06-23 00:14:41 -04:00
parent 6f6aec2930
commit 6483335df0
2 changed files with 76 additions and 0 deletions

View file

@ -47,6 +47,13 @@ _MAX_TOKEN_SUPPORT_MODES: frozenset[str] = frozenset(
{"chat", "completion", "responses"}
)
# Tool config is meaningless on a trivial health probe and makes some providers
# fail it outright (OpenRouter server tools like `openrouter:web_search` 500 the
# probe). Stripped from both the top-level params and `extra_body` before probing.
_HEALTH_CHECK_STRIPPED_REQUEST_KEYS: frozenset[str] = frozenset(
{"tools", "tool_choice"}
)
def _resolve_health_check_mode(
model_info: Mapping[str, object], litellm_params: Mapping[str, object]
@ -452,7 +459,20 @@ def _update_litellm_params_for_health_check(
- updates the `model` param with the `health_check_model` if it exists Doc: https://docs.litellm.ai/docs/proxy/health#wildcard-routes
- updates the `voice` param with the `health_check_voice` for `audio_speech` mode if it exists Doc: https://docs.litellm.ai/docs/proxy/health#text-to-speech-models
- for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID, and pins `custom_llm_provider` to `bedrock` (only when the deployment hasn't already set one, so an explicit `bedrock_converse` survives) so the bare model id still resolves to the provider (e.g. cross-region ids like `us.cohere.embed-v4:0`)
- works on a copy and strips `tools`/`tool_choice` from the top level and `extra_body`; a trivial probe never needs tool config and some providers (e.g. OpenRouter server tools) 500 when it is present
"""
litellm_params = {
k: v
for k, v in litellm_params.items()
if k not in _HEALTH_CHECK_STRIPPED_REQUEST_KEYS
}
_extra_body = litellm_params.get("extra_body")
if isinstance(_extra_body, dict):
litellm_params["extra_body"] = {
k: v
for k, v in _extra_body.items()
if k not in _HEALTH_CHECK_STRIPPED_REQUEST_KEYS
}
mode = _resolve_health_check_mode(
model_info, litellm_params # any-ok: untyped router config dict
)

View file

@ -447,6 +447,62 @@ def test_update_litellm_params_for_health_check():
)
def test_health_check_strips_tool_config_from_probe():
"""
Regression for #31008: health probes must not carry tool configuration.
OpenRouter server tools configured via `extra_body.tools` (e.g.
`openrouter:web_search`) made the trivial probe 500. `tools`/`tool_choice`
are stripped from both the top level and `extra_body`, while other
`extra_body` fields (e.g. provider routing) are preserved, and the caller's
params are never mutated.
"""
from litellm.proxy.health_check import _update_litellm_params_for_health_check
extra_body = {
"tools": [{"type": "openrouter:web_search"}],
"tool_choice": "auto",
"provider": {"order": ["OpenAI"]},
}
litellm_params = {
"model": "openrouter/openai/gpt-4o-mini",
"api_key": "fake_key",
"tools": [{"type": "function", "function": {"name": "f"}}],
"tool_choice": "auto",
"extra_body": extra_body,
}
updated = _update_litellm_params_for_health_check({}, litellm_params)
# probe is stripped of all tool config, top-level and nested
assert "tools" not in updated and "tool_choice" not in updated
assert "tools" not in updated["extra_body"]
assert "tool_choice" not in updated["extra_body"]
assert updated["extra_body"]["provider"] == {"order": ["OpenAI"]}
assert isinstance(updated["messages"], list)
# caller's params are not mutated: a fresh object, no probe fields leaked in,
# tool config intact both top-level and nested
assert updated is not litellm_params
assert "messages" not in litellm_params
assert "tools" in litellm_params and "tool_choice" in litellm_params
assert extra_body["tools"] == [{"type": "openrouter:web_search"}]
assert extra_body["tool_choice"] == "auto"
assert updated["extra_body"] is not extra_body
def test_health_check_probe_without_extra_body_does_not_crash():
"""A deployment with no `extra_body` must still build a valid probe."""
from litellm.proxy.health_check import _update_litellm_params_for_health_check
updated = _update_litellm_params_for_health_check(
{}, {"model": "openai/gpt-4o-mini", "api_key": "fake_key"}
)
assert "extra_body" not in updated
assert isinstance(updated["messages"], list)
@pytest.mark.asyncio
async def test_perform_health_check_filters_by_model_id():
"""