mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-01 02:02:20 +00:00
fix(compact_20260112): enforce parent key/team allowlist on summary model
The compact_20260112 polyfill summary subrequest used llm_router.acompletion directly, bypassing the proxy auth checks that gate model access for the parent key/team. A caller whose key/team was not authorized for the configured context_management_summary_model could still cause the proxy to invoke that model and return its output as a compaction block. Pull the parent's UserAPIKeyAuth out of litellm_metadata in the handler, thread it through the dispatcher into apply_compact_20260112, and gate the summary call on _can_object_call_model for both key-level and team-level allowlists. Failures land as applied_edits[0].error = summary_model_access_denied without raising. SDK callers (no UserAPIKeyAuth) remain unaffected.
This commit is contained in:
mateo-berri
parent
b989b6fedf
commit
488d9eebb2
4 changed files with 248 additions and 1 deletions
|
|
@ -52,6 +52,23 @@ def _messages_have_compaction_block(messages: List[Dict]) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _extract_user_api_key_auth(kwargs: Dict[str, Any]) -> Any:
|
||||
"""Pull the parent request's ``UserAPIKeyAuth`` out of ``litellm_metadata``.
|
||||
|
||||
The proxy attaches the full auth object under
|
||||
``data["litellm_metadata"]["user_api_key_auth"]`` (see
|
||||
``LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata``).
|
||||
The context-management polyfill uses it to gate the summary subrequest on
|
||||
the parent key/team's model allowlist; without it, the summary call would
|
||||
bypass the proxy auth checks. Returns ``None`` for SDK callers that bypass
|
||||
the proxy entirely.
|
||||
"""
|
||||
litellm_metadata = kwargs.get("litellm_metadata")
|
||||
if not isinstance(litellm_metadata, dict):
|
||||
return None
|
||||
return litellm_metadata.get("user_api_key_auth")
|
||||
|
||||
|
||||
async def _prepare_context_managed_request(
|
||||
*,
|
||||
model: str,
|
||||
|
|
@ -62,6 +79,7 @@ async def _prepare_context_managed_request(
|
|||
metadata: Optional[Dict],
|
||||
drop_params: Optional[bool],
|
||||
llm_router: Any,
|
||||
user_api_key_auth: Any = None,
|
||||
) -> Optional[PolyfillResult]:
|
||||
"""Apply client compaction history, then optional context_management polyfill."""
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
|
||||
|
|
@ -102,6 +120,7 @@ async def _prepare_context_managed_request(
|
|||
metadata=metadata,
|
||||
drop_params=drop_params,
|
||||
llm_router=llm_router,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
|
||||
if polyfill_result is not None:
|
||||
|
|
@ -221,6 +240,7 @@ async def _run_polyfill_if_enabled(
|
|||
metadata: Optional[Dict],
|
||||
drop_params: Optional[bool],
|
||||
llm_router: Any,
|
||||
user_api_key_auth: Any = None,
|
||||
) -> Optional[PolyfillResult]:
|
||||
"""Run the async context_management polyfill if a spec is present.
|
||||
|
||||
|
|
@ -247,6 +267,7 @@ async def _run_polyfill_if_enabled(
|
|||
context_management_spec=context_management_spec,
|
||||
metadata=metadata,
|
||||
llm_router=llm_router,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
except AnthropicContextManagementError:
|
||||
# Surface validation errors so the endpoint can emit an Anthropic-format
|
||||
|
|
@ -557,6 +578,8 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
user_api_key_auth = _extract_user_api_key_auth(kwargs)
|
||||
|
||||
polyfill_result = await _prepare_context_managed_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
@ -566,6 +589,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
metadata=metadata,
|
||||
drop_params=drop_params,
|
||||
llm_router=litellm_router,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
|
||||
effective_messages = (
|
||||
|
|
@ -708,6 +732,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
metadata=metadata,
|
||||
drop_params=drop_params,
|
||||
llm_router=litellm_router,
|
||||
user_api_key_auth=_extract_user_api_key_auth(kwargs),
|
||||
)
|
||||
|
||||
effective_messages = (
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ async def apply_context_management(
|
|||
context_management_spec: Union[Dict[str, Any], List[Dict[str, Any]], None],
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
llm_router: Any = None,
|
||||
user_api_key_auth: Any = None,
|
||||
) -> PolyfillResult:
|
||||
"""Run edits in order; return a single ``PolyfillResult``.
|
||||
|
||||
|
|
@ -102,6 +103,7 @@ async def apply_context_management(
|
|||
if inspect.iscoroutinefunction(editor):
|
||||
kwargs["metadata"] = metadata
|
||||
kwargs["llm_router"] = llm_router
|
||||
kwargs["user_api_key_auth"] = user_api_key_auth
|
||||
raw_result = await cast(Callable[..., Awaitable[Any]], editor)(**kwargs)
|
||||
else:
|
||||
raw_result = editor(**kwargs)
|
||||
|
|
|
|||
|
|
@ -65,6 +65,63 @@ def _read_summary_model_setting() -> Optional[str]:
|
|||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _check_summary_model_access(
|
||||
user_api_key_auth: Any,
|
||||
summary_model: str,
|
||||
llm_router: Any,
|
||||
) -> bool:
|
||||
"""Return True when the parent key/team is authorized to call ``summary_model``.
|
||||
|
||||
The summary subrequest does not pass through ``user_api_key_auth`` again,
|
||||
so without this gate a caller whose key/team is not allowed to use
|
||||
``context_management_summary_model`` could still get the proxy to invoke
|
||||
that model and return its ``<summary>`` output as a compaction block.
|
||||
|
||||
Returns True (allow) when ``user_api_key_auth`` is not present — SDK
|
||||
callers and tests run outside the proxy, where no key/team policy exists.
|
||||
Returns False when either the key-level or team-level allowlist denies
|
||||
the summary model.
|
||||
"""
|
||||
if user_api_key_auth is None:
|
||||
return True
|
||||
try:
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
key_models = list(getattr(user_api_key_auth, "models", None) or [])
|
||||
team_id = getattr(user_api_key_auth, "team_id", None)
|
||||
team_model_aliases = getattr(user_api_key_auth, "team_model_aliases", None)
|
||||
if key_models:
|
||||
try:
|
||||
_can_object_call_model(
|
||||
model=summary_model,
|
||||
llm_router=llm_router,
|
||||
models=key_models,
|
||||
team_model_aliases=team_model_aliases,
|
||||
team_id=team_id,
|
||||
object_type="key",
|
||||
)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
team_models = list(getattr(user_api_key_auth, "team_models", None) or [])
|
||||
if team_models:
|
||||
try:
|
||||
_can_object_call_model(
|
||||
model=summary_model,
|
||||
llm_router=llm_router,
|
||||
models=team_models,
|
||||
team_model_aliases=team_model_aliases,
|
||||
team_id=team_id,
|
||||
object_type="team",
|
||||
)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _find_latest_compaction_index(
|
||||
messages: List[Dict[str, Any]],
|
||||
) -> Tuple[Optional[int], Optional[int]]:
|
||||
|
|
@ -544,7 +601,7 @@ def apply_client_compaction_block_history(
|
|||
)
|
||||
|
||||
|
||||
async def apply_compact_20260112(
|
||||
async def apply_compact_20260112( # noqa: PLR0915
|
||||
*,
|
||||
model: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
|
|
@ -553,6 +610,7 @@ async def apply_compact_20260112(
|
|||
edit_spec: Dict[str, Any],
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
llm_router: Any = None,
|
||||
user_api_key_auth: Any = None,
|
||||
) -> PolyfillResult:
|
||||
"""Apply ``compact_20260112``; return a ``PolyfillResult``.
|
||||
|
||||
|
|
@ -652,6 +710,23 @@ async def apply_compact_20260112(
|
|||
# Phase C: summarize. ``augmented_system`` carries any prior compaction
|
||||
# summary so multi-round compaction does not lose accumulated history —
|
||||
# ``effective_messages`` only contains turns since the last compaction.
|
||||
if not _check_summary_model_access(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
summary_model=summary_model,
|
||||
llm_router=llm_router,
|
||||
):
|
||||
verbose_logger.warning(
|
||||
"compact_20260112: caller not authorized for summary_model=%s; "
|
||||
"skipping summary call",
|
||||
summary_model,
|
||||
)
|
||||
applied["error"] = "summary_model_access_denied"
|
||||
return PolyfillResult(
|
||||
messages=downstream_messages,
|
||||
system=augmented_system,
|
||||
applied_edits=[applied],
|
||||
)
|
||||
|
||||
prompt = _build_summary_prompt(edit_spec, tools)
|
||||
summary_messages = _build_summary_messages(
|
||||
effective_messages, prompt, system=augmented_system
|
||||
|
|
|
|||
|
|
@ -1049,6 +1049,151 @@ async def test_summary_call_sends_default_max_tokens():
|
|||
assert captured_kwargs.get("max_tokens") == COMPACT_SUMMARY_MAX_TOKENS
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Editor: summary model key/team access gate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _fake_user_api_key_auth(*, key_models=None, team_models=None, team_id=None):
|
||||
"""Build a minimal stand-in for ``UserAPIKeyAuth`` with just the fields
|
||||
consulted by ``_check_summary_model_access``. Avoids pulling the proxy
|
||||
deps into this unit test."""
|
||||
|
||||
class _Auth:
|
||||
pass
|
||||
|
||||
auth = _Auth()
|
||||
auth.models = list(key_models) if key_models is not None else []
|
||||
auth.team_models = list(team_models) if team_models is not None else []
|
||||
auth.team_id = team_id
|
||||
auth.team_model_aliases = None
|
||||
return auth
|
||||
|
||||
|
||||
async def test_summary_model_denied_when_key_not_in_allowlist():
|
||||
"""Caller key restricted to specific models cannot trigger an unauthorized summary model."""
|
||||
messages = _simple_messages()
|
||||
mock_call = AsyncMock(return_value=_make_mock_response("<summary>x</summary>"))
|
||||
|
||||
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,
|
||||
),
|
||||
):
|
||||
result = await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=_fake_user_api_key_auth(key_models=["gpt-4o"]),
|
||||
)
|
||||
|
||||
mock_call.assert_not_awaited()
|
||||
assert result.compaction_block is None
|
||||
assert result.iterations_usage is None
|
||||
assert result.applied_edits[0]["type"] == "compact_20260112"
|
||||
assert result.applied_edits[0].get("error") == "summary_model_access_denied"
|
||||
|
||||
|
||||
async def test_summary_model_denied_when_team_not_in_allowlist():
|
||||
"""Team-level model allowlist is enforced even if the key allows all models."""
|
||||
messages = _simple_messages()
|
||||
mock_call = AsyncMock(return_value=_make_mock_response("<summary>x</summary>"))
|
||||
|
||||
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,
|
||||
),
|
||||
):
|
||||
result = await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=_fake_user_api_key_auth(
|
||||
key_models=["all-proxy-models"], team_models=["gpt-4o"]
|
||||
),
|
||||
)
|
||||
|
||||
mock_call.assert_not_awaited()
|
||||
assert result.applied_edits[0].get("error") == "summary_model_access_denied"
|
||||
|
||||
|
||||
async def test_summary_model_allowed_when_in_key_allowlist():
|
||||
"""Caller key that explicitly allows the summary model is permitted to use it."""
|
||||
messages = _simple_messages()
|
||||
mock_call = AsyncMock(return_value=_make_mock_response("<summary>ok</summary>"))
|
||||
|
||||
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,
|
||||
),
|
||||
):
|
||||
result = await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=_fake_user_api_key_auth(
|
||||
key_models=["claude-haiku-4-5", "gpt-4o"]
|
||||
),
|
||||
)
|
||||
|
||||
mock_call.assert_awaited_once()
|
||||
assert result.compaction_block is not None
|
||||
assert result.compaction_block["content"] == "ok"
|
||||
assert not result.applied_edits[0].get("error")
|
||||
|
||||
|
||||
async def test_summary_model_allowed_when_no_user_api_key_auth():
|
||||
"""SDK callers (no proxy auth object) are not gated."""
|
||||
messages = _simple_messages()
|
||||
mock_call = AsyncMock(return_value=_make_mock_response("<summary>ok</summary>"))
|
||||
|
||||
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,
|
||||
),
|
||||
):
|
||||
result = await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
)
|
||||
|
||||
mock_call.assert_awaited_once()
|
||||
assert result.compaction_block is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dispatcher integration: compact_20260112 via apply_context_management
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue