diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index ea77fd07b34..da286c9482d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -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 = ( diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py index 0c2a13869e8..5b8f1c267f4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -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) diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index ef404c45eeb..5959ed38655 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -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 ```` 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 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 059414ea695..a54e2146fb0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -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("x")) + + 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("x")) + + 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("ok")) + + 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("ok")) + + 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 # ---------------------------------------------------------------------------