fix(compact_20260112): propagate allowed_model_region to summary subrequest

The router enforces region restrictions by reading allowed_model_region
from top-level request kwargs (Router._common_checks_available_deployment),
but the compact_20260112 summary subrequest only forwarded litellm_metadata.
A region-restricted caller could trigger compaction and have their conversation
summarized by a deployment outside the permitted region.

Extract allowed_model_region from user_api_key_auth and pass it through
_call_summary_model as a top-level kwarg so the router applies the same
region constraints the parent request would.
This commit is contained in:
mateo-berri 2026-05-28 07:24:03 +00:00
parent 0adbb24633
commit d41ab61409
No known key found for this signature in database
2 changed files with 94 additions and 0 deletions

View file

@ -525,6 +525,7 @@ async def _call_summary_model(
summary_messages: List[Dict[str, Any]],
metadata: Dict[str, Any],
llm_router: Any,
allowed_model_region: Optional[str] = None,
) -> Any:
"""Invoke the configured summary model.
@ -541,12 +542,19 @@ async def _call_summary_model(
# post-call spend hooks read for budget attribution. The provider-level
# ``metadata`` kwarg corresponds to the upstream API request body and would
# not flow into spend tracking.
# ``allowed_model_region`` must travel as a top-level kwarg because the
# router enforces region restrictions by reading ``request_kwargs`` directly
# (see ``Router._common_checks_available_deployment``); without this the
# summary subrequest could be routed to a deployment outside the caller's
# permitted region.
call_kwargs: Dict[str, Any] = {
"model": summary_model,
"messages": summary_messages,
"max_tokens": COMPACT_SUMMARY_MAX_TOKENS,
"litellm_metadata": metadata,
}
if allowed_model_region is not None:
call_kwargs["allowed_model_region"] = allowed_model_region
if llm_router is not None and hasattr(llm_router, "acompletion"):
return await llm_router.acompletion(**call_kwargs)
return await litellm.acompletion(**call_kwargs)
@ -754,6 +762,7 @@ async def apply_compact_20260112( # noqa: PLR0915
effective_messages, prompt, system=augmented_system
)
propagated_metadata = _propagate_metadata(litellm_metadata)
allowed_model_region = getattr(user_api_key_auth, "allowed_model_region", None)
try:
response = await _call_summary_model(
@ -761,6 +770,7 @@ async def apply_compact_20260112( # noqa: PLR0915
summary_messages=summary_messages,
metadata=propagated_metadata,
llm_router=llm_router,
allowed_model_region=allowed_model_region,
)
except Exception as e:
verbose_logger.warning("compact_20260112: summary call failed: %s", e)

View file

@ -1194,6 +1194,90 @@ async def test_summary_model_allowed_when_no_user_api_key_auth():
assert result.compaction_block is not None
async def test_summary_call_propagates_allowed_model_region():
"""``allowed_model_region`` from ``user_api_key_auth`` is propagated to the
summary subrequest as a top-level kwarg so the router applies the same
region restriction the parent request would.
"""
messages = _simple_messages()
mock_call = AsyncMock(return_value=_make_mock_response("<summary>ok</summary>"))
auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
auth.allowed_model_region = "eu"
with (
patch(
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
return_value="claude-haiku-4-5",
),
patch("litellm.token_counter", return_value=200_000),
patch(
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
mock_call,
),
):
await apply_compact_20260112(
model=MODEL,
messages=messages,
tools=None,
system=None,
edit_spec=_EDIT_SPEC_DEFAULT,
user_api_key_auth=auth,
)
mock_call.assert_awaited_once()
assert mock_call.await_args.kwargs.get("allowed_model_region") == "eu"
async def test_summary_call_omits_allowed_model_region_when_unset():
"""Callers without a region restriction must not get an ``allowed_model_region=None``
kwarg, which would otherwise force the router to evaluate region filtering.
"""
from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
_call_summary_model,
)
captured_kwargs: dict = {}
class _FakeRouter:
async def acompletion(self, **kwargs):
captured_kwargs.update(kwargs)
return _make_mock_response("<summary>x</summary>")
await _call_summary_model(
summary_model="claude-haiku-4-5",
summary_messages=[{"role": "user", "content": "hi"}],
metadata={},
llm_router=_FakeRouter(),
)
assert "allowed_model_region" not in captured_kwargs
async def test_summary_call_forwards_allowed_model_region_when_set():
"""When the caller is region-restricted, the kwarg reaches the router."""
from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
_call_summary_model,
)
captured_kwargs: dict = {}
class _FakeRouter:
async def acompletion(self, **kwargs):
captured_kwargs.update(kwargs)
return _make_mock_response("<summary>x</summary>")
await _call_summary_model(
summary_model="claude-haiku-4-5",
summary_messages=[{"role": "user", "content": "hi"}],
metadata={},
llm_router=_FakeRouter(),
allowed_model_region="eu",
)
assert captured_kwargs.get("allowed_model_region") == "eu"
# ---------------------------------------------------------------------------
# Dispatcher integration: compact_20260112 via apply_context_management
# ---------------------------------------------------------------------------