From fc344fd2a93595466e31f5ee08d3e2341deec5f6 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 22:00:52 -0700 Subject: [PATCH] chore(release): patch v1.89.0-rc.2 with batch-file auth, CrowdStrike AIDR, Mantle SigV4, NetApp streaming-cost fix, and team-scoped Datadog toward v1.89.0-rc.3 (#30179) * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) (#30009) * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead. Co-authored-by: Cursor * fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593) Restores the reverse-lookup for the JSONL body.model fallback path so that legacy/pre-target_model_names managed files still map stripped provider IDs back to proxy aliases before auth. Also cleans up redundant `or None`. Co-Authored-By: Claude Sonnet 4.6 * Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)" This reverts commit 30d2e96f77ef521ccaaf2193fe554980380eb669. --------- Co-authored-by: Cursor Co-authored-by: Claude Sonnet 4.6 (cherry picked from commit 2cd7e874859eade595928b97b73fab5e10620629) * feat(guardrails): capture user and model metadata in CrowdStrike AIDR (cherry picked from commit 6fc715c5bd5b9acd1c73e755da8ecbe808765dfc) * fix(guardrails): read CrowdStrike AIDR identity from both metadata bags (#29991) Capture user_id and extra_info from metadata or litellm_metadata. The single-bag read dropped identity whenever a request carried a present litellm_metadata field (null or a user-supplied dict), since /chat/completions routes the authenticated identity into metadata while the guardrail read litellm_metadata first (cherry picked from commit 1bbaf1c39dda367f5a2b4b6b9ab4cac46d71ab14) * feat(bedrock_mantle): add SigV4/IAM auth to Responses API route (#29788) Applied as the squash diff of PR #29788 (head 9800b2f17c), which landed upstream inside the litellm_oss_staging_080626 sync (32c88ca74f, #29932) and has no standalone commit to cherry-pick. The rc line already carries the prerequisite #29490 Responses route via the 040626 sync. * fix: completion_cost AttributeError on streaming Anthropic web_search responses (#26153) (#27346) Cherry-picked from staging squash 4a3860df1f. The rc line predates the Usage.__init__ server_tool_use dict->ServerToolUse coercion that staging carries (it landed via the squashed OSS sync #29932 / 32c88ca74f, not as a standalone commit). The calculate_usage Usage(**returned_usage.model_dump()) round-trip re-serializes server_tool_use to a plain dict, so without that coercion the rebuilt usage holds a dict and the regression test asserting a ServerToolUse type fails. Restored the coercion in litellm/types/utils.py to satisfy the prerequisite -- it matches #27346's own first commit (coerce server_tool_use dict to ServerToolUse in Usage.__init__), which was dropped from the squash only because staging already carried it. * feat(datadog): add team-scoped Datadog callback support (#29947) Cherry-picked from the PR head 9c049daa1b (single-commit PR, merged to litellm_oss_branch). Applied cleanly; no conflicts. Note: black --check in this worktree flags pre-existing multi-line string formatting in litellm_core_utils/litellm_logging.py (lines ~1006-1050) that is already present on the patch/v1.89.0-rc.1 base and is untouched by this pick -- left as-is to avoid reformatting unrelated lines. --------- Co-authored-by: Sameer Kankute Co-authored-by: Cursor Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Kenan Yildirim Co-authored-by: yuneng-jiang Co-authored-by: Kent Co-authored-by: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Co-authored-by: aanchal22 <12680748+aanchal22@users.noreply.github.com> --- litellm/integrations/datadog/datadog.py | 69 ++- .../datadog/datadog_team_handler.py | 117 +++++ .../initialize_dynamic_callback_params.py | 8 + litellm/litellm_core_utils/litellm_logging.py | 40 +- .../llm_cost_calc/tool_call_cost_tracking.py | 7 +- .../litellm_core_utils/llm_cost_calc/utils.py | 22 +- .../streaming_chunk_builder_utils.py | 13 +- litellm/llms/anthropic/cost_calculation.py | 14 +- .../llms/base_llm/responses/transformation.py | 20 + .../responses/transformation.py | 120 +++++- litellm/llms/custom_httpx/llm_http_handler.py | 105 +++-- .../crowdstrike_aidr/crowdstrike_aidr.py | 28 +- litellm/proxy/hooks/batch_rate_limiter.py | 29 +- litellm/types/utils.py | 11 +- .../datadog/test_datadog_team_handler.py | 194 +++++++++ ...est_tool_call_cost_tracking_dict_safety.py | 88 ++++ ...streaming_chunk_builder_server_tool_use.py | 130 ++++++ .../test_streaming_chunk_builder_utils.py | 5 +- .../test_cost_calculation_dict_safety.py | 94 +++++ ...bedrock_mantle_responses_transformation.py | 398 +++++++++++++++++- .../custom_httpx/test_llm_http_handler.py | 238 +++++++++++ .../guardrail_hooks/test_crowdstrike_aidr.py | 194 +++++++++ .../proxy/hooks/test_batch_file_validation.py | 99 ++++- 23 files changed, 1948 insertions(+), 95 deletions(-) create mode 100644 litellm/integrations/datadog/datadog_team_handler.py create mode 100644 tests/test_litellm/integrations/datadog/test_datadog_team_handler.py create mode 100644 tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py create mode 100644 tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py create mode 100644 tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 79a9219a39c..a79a852ff29 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -92,12 +92,22 @@ class DataDogLogger( # Class variables or attributes def __init__( self, + dd_api_key: Optional[str] = None, + dd_site: Optional[str] = None, + dd_agent_host: Optional[str] = None, + dd_agent_port: Optional[str] = None, **kwargs, ): """ Initializes the datadog logger, checks if the correct env variables are set - Required environment variables (Direct API): + Args: + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. + dd_site: Datadog site (e.g. "us5.datadoghq.com"). Falls back to DD_SITE env var. + dd_agent_host: Hostname or IP of DataDog agent. Falls back to LITELLM_DD_AGENT_HOST env var. + dd_agent_port: Port of DataDog agent (default: 10518). Falls back to LITELLM_DD_AGENT_PORT env var. + + Required environment variables (Direct API) when kwargs not provided: `DD_API_KEY` - your datadog api key `DD_SITE` - your datadog site, example = `"us5.datadoghq.com"` @@ -130,12 +140,19 @@ class DataDogLogger( ) # Configure DataDog endpoint (Agent or Direct API) - # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST - dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") - if dd_agent_host: - self._configure_dd_agent(dd_agent_host=dd_agent_host) + # Prefer explicit kwargs, then fall back to env vars + resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST") + if resolved_agent_host: + self._configure_dd_agent( + dd_agent_host=resolved_agent_host, + dd_agent_port=dd_agent_port, + dd_api_key=dd_api_key, + ) else: - self._configure_dd_direct_api() + self._configure_dd_direct_api( + dd_api_key=dd_api_key, + dd_site=dd_site, + ) # Optional override for testing dd_base_url = get_datadog_base_url_from_env() @@ -172,34 +189,54 @@ class DataDogLogger( ).model_dump() return dict_datadog_params - def _configure_dd_agent(self, dd_agent_host: str) -> None: + def _configure_dd_agent( + self, + dd_agent_host: str, + dd_agent_port: Optional[str] = None, + dd_api_key: Optional[str] = None, + ) -> None: """ Configure DataDog Agent for log forwarding Args: dd_agent_host: Hostname or IP of DataDog agent + dd_agent_port: Port of DataDog agent. Falls back to LITELLM_DD_AGENT_PORT env var (default: 10518). + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. Optional when using agent. """ - dd_agent_port = os.getenv( + resolved_port = dd_agent_port or os.getenv( "LITELLM_DD_AGENT_PORT", "10518" ) # default port for logs - self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs" - self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent + self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs" + self.DD_API_KEY = dd_api_key or os.getenv( + "DD_API_KEY" + ) # Optional when using agent verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}") - def _configure_dd_direct_api(self) -> None: + def _configure_dd_direct_api( + self, + dd_api_key: Optional[str] = None, + dd_site: Optional[str] = None, + ) -> None: """ Configure direct DataDog API connection + Args: + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. + dd_site: Datadog site. Falls back to DD_SITE env var. + Raises: - Exception: If required environment variables are not set + Exception: If required credentials are not provided via args or env vars """ - if os.getenv("DD_API_KEY", None) is None: + resolved_api_key = dd_api_key or os.getenv("DD_API_KEY") + resolved_site = dd_site or os.getenv("DD_SITE") + + if resolved_api_key is None: raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>") - if os.getenv("DD_SITE", None) is None: + if resolved_site is None: raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>") - self.DD_API_KEY = os.getenv("DD_API_KEY") - self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs" + self.DD_API_KEY = resolved_api_key + self.intake_url = f"https://http-intake.logs.{resolved_site}/api/v2/logs" async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ diff --git a/litellm/integrations/datadog/datadog_team_handler.py b/litellm/integrations/datadog/datadog_team_handler.py new file mode 100644 index 00000000000..dda24c8e6a2 --- /dev/null +++ b/litellm/integrations/datadog/datadog_team_handler.py @@ -0,0 +1,117 @@ +""" +DataDog Team Handler + +Used to get the DataDogLogger for a given request. +Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams + +from .datadog import DataDogLogger + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache +else: + DynamicLoggingCache = Any + + +class DatadogLoggingConfig(TypedDict): + dd_api_key: Optional[str] + dd_site: Optional[str] + dd_agent_host: Optional[str] + dd_agent_port: Optional[str] + + +class DataDogHandler: + @staticmethod + def get_datadog_logger_for_request( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + in_memory_dynamic_logger_cache: DynamicLoggingCache, + ) -> DataDogLogger: + """ + Get a team-scoped DataDogLogger for a given request. + + Resolves and caches per-team DataDogLogger instances using DynamicLoggingCache, + keyed by the team's DD credentials. Each unique set of credentials gets its own + logger instance with its own batch/flush loop. + + Note: This handler is only called when team-scoped DD credentials are present. + The global (env-var based) DataDogLogger is managed separately by + _init_custom_logger_compatible_class via _in_memory_loggers. + """ + _credentials = DataDogHandler.get_dynamic_datadog_logging_config( + standard_callback_dynamic_params=standard_callback_dynamic_params, + ) + credentials_dict = dict(_credentials) + + # check if datadog logger is already cached + temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache( + credentials=credentials_dict, service_name="datadog" + ) + + # if not cached, create a new datadog logger and cache it + if temp_datadog_logger is None: + temp_datadog_logger = ( + DataDogHandler._create_datadog_logger_from_credentials( + credentials=credentials_dict, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + ) + + return temp_datadog_logger + + @staticmethod + def _create_datadog_logger_from_credentials( + credentials: Dict, + in_memory_dynamic_logger_cache: DynamicLoggingCache, + ) -> DataDogLogger: + """ + Create a DataDogLogger from the credentials and cache it. + """ + datadog_logger = DataDogLogger( + dd_api_key=credentials.get("dd_api_key"), + dd_site=credentials.get("dd_site"), + dd_agent_host=credentials.get("dd_agent_host"), + dd_agent_port=credentials.get("dd_agent_port"), + ) + in_memory_dynamic_logger_cache.set_cache( + credentials=credentials, + service_name="datadog", + logging_obj=datadog_logger, + ) + verbose_logger.debug( + "Datadog: Created and cached new DataDogLogger for team-scoped credentials" + ) + return datadog_logger + + @staticmethod + def get_dynamic_datadog_logging_config( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> DatadogLoggingConfig: + """ + Get the Datadog logging config for a given request from dynamic params. + """ + return DatadogLoggingConfig( + dd_api_key=standard_callback_dynamic_params.get("dd_api_key"), + dd_site=standard_callback_dynamic_params.get("dd_site"), + dd_agent_host=standard_callback_dynamic_params.get("dd_agent_host"), + dd_agent_port=standard_callback_dynamic_params.get("dd_agent_port"), + ) + + @staticmethod + def _dynamic_datadog_credentials_are_passed( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> bool: + """ + Check if dynamic Datadog credentials are passed in standard_callback_dynamic_params. + """ + if ( + standard_callback_dynamic_params.get("dd_api_key") is not None + or standard_callback_dynamic_params.get("dd_site") is not None + or standard_callback_dynamic_params.get("dd_agent_host") is not None + ): + return True + return False diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index a89dae52316..949076aabf3 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -53,11 +53,19 @@ _supported_callback_params = [ "braintrust_host", "slack_webhook_url", "lunary_public_key", + "dd_api_key", + "dd_site", + "dd_agent_host", + "dd_agent_port", ] _request_blocked_callback_params = { "gcs_bucket_name", "gcs_path_service_account", + "dd_api_key", + "dd_site", + "dd_agent_host", + "dd_agent_port", } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f20b66790c4..a76ca954670 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -376,13 +376,14 @@ class Logging(LiteLLMLoggingBaseClass): List[Union[str, Callable, CustomLogger]] ] = dynamic_async_failure_callbacks - # Process dynamic callbacks - self.process_dynamic_callbacks() - ## DYNAMIC LANGFUSE / GCS / logging callback KEYS ## self.standard_callback_dynamic_params: StandardCallbackDynamicParams = ( self.initialize_standard_callback_dynamic_params(kwargs) ) + + # Process dynamic callbacks (after standard_callback_dynamic_params is initialized, + # so team-scoped credentials are available for callback initialization) + self.process_dynamic_callbacks() self.standard_built_in_tools_params: StandardBuiltInToolsParams = ( self.initialize_standard_built_in_tools_params(kwargs) ) @@ -477,8 +478,21 @@ class Logging(LiteLLMLoggingBaseClass): isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks ): + # For callbacks that support team-scoped credentials (e.g. datadog), + # pass only the relevant dynamic params as custom_logger_init_args. + _custom_logger_init_args: Optional[dict] = None + if callback == "datadog": + _custom_logger_init_args = { + k: v + for k, v in self.standard_callback_dynamic_params.items() + if k.startswith("dd_") + } + callback_class = _init_custom_logger_compatible_class( - callback, internal_usage_cache=None, llm_router=None # type: ignore + callback, # type: ignore[arg-type] + internal_usage_cache=None, + llm_router=None, # type: ignore + custom_logger_init_args=_custom_logger_init_args, ) if callback_class is not None: processed_list.append(callback_class) @@ -3890,6 +3904,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(_prometheus_logger) return _prometheus_logger # type: ignore elif logging_integration == "datadog": + # Check if team-scoped credentials are provided + _dd_api_key = custom_logger_init_args.get("dd_api_key") + _dd_site = custom_logger_init_args.get("dd_site") + _dd_agent_host = custom_logger_init_args.get("dd_agent_host") + _dd_agent_port = custom_logger_init_args.get("dd_agent_port") + + if _dd_api_key or _dd_site or _dd_agent_host: + # Team-scoped credentials: use DynamicLoggingCache for per-credential isolation + from litellm.integrations.datadog.datadog_team_handler import ( + DataDogHandler, + ) + + return DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=custom_logger_init_args, # type: ignore + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + + # Global (env-var based): reuse cached instance for callback in _in_memory_loggers: if isinstance(callback, DataDogLogger): return callback # type: ignore diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 8da66d4600d..413ddb71bf8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -6,6 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, @@ -339,8 +340,7 @@ class StandardBuiltInToolCostTracking: # and _handle_web_search_cost() is never called. if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True return False @@ -352,8 +352,7 @@ class StandardBuiltInToolCostTracking: elif usage is not None: if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True elif ( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f39c942f90f..93049adf75a 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1,7 +1,7 @@ # What is this? ## Helper utilities for cost_per_token() -from typing import Literal, Optional, Tuple, TypedDict, cast +from typing import Any, Literal, Optional, Tuple, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]: return value if isinstance(value, int) else None +def _get_web_search_requests(server_tool_use: Any) -> Optional[int]: + """ + Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value + that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, + or any other object supporting attribute access. + + Returns ``None`` when the value cannot be resolved — callers can + distinguish "absent" from "zero" using ``is None``. + + See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder`` + historically left this as a plain ``dict``, which broke direct attribute + access in cost calculation. + """ + if server_tool_use is None: + return None + if isinstance(server_tool_use, dict): + return server_tool_use.get("web_search_requests") + return getattr(server_tool_use, "web_search_requests", None) + + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: return True diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c3842..c2a17ae8dcc 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -588,7 +588,18 @@ class ChunkProcessor: hasattr(usage_chunk, "server_tool_use") and usage_chunk.server_tool_use is not None ): - server_tool_use = usage_chunk.server_tool_use + # Coerce dict to ServerToolUse so downstream cost-calc code + # (which accesses .web_search_requests as an attribute) + # doesn't raise AttributeError. Some providers / streaming + # paths leave server_tool_use as a plain dict on the chunk. + if isinstance(usage_chunk.server_tool_use, dict): + server_tool_use = ServerToolUse(**usage_chunk.server_tool_use) + elif isinstance(usage_chunk.server_tool_use, ServerToolUse): + server_tool_use = usage_chunk.server_tool_use + else: + server_tool_use = ServerToolUse.model_validate( + usage_chunk.server_tool_use + ) if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 3882d8f978c..6a031498dae 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional, Tuple from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_token_base_cost, + _get_web_search_requests, _parse_prompt_tokens_details, calculate_cache_writing_cost, generic_cost_per_token, @@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search( if model_info is None: return 0.0 - if ( - usage is None - or usage.server_tool_use is None - or usage.server_tool_use.web_search_requests is None - ): + if usage is None: + return 0.0 + web_search_requests = _get_web_search_requests( + getattr(usage, "server_tool_use", None) + ) + if web_search_requests is None: return 0.0 ## Get the cost per web search request @@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search( return 0.0 ## Calculate the total cost - total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests + total_cost = cost_per_web_search_request * web_search_requests return total_cost diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 853eb282758..407d5ad8146 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -62,6 +62,26 @@ class BaseResponsesAPIConfig(ABC): """ return False + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + """Sign the request after the body is finalized. + + Default is a no-op (returns headers unchanged, no signed body). Providers + whose endpoint requires request signing (e.g. Bedrock Mantle SigV4) + override this and return the signed body bytes so the handler sends those + exact bytes. + """ + return headers, None + @abstractmethod def get_supported_openai_params(self, model: str) -> list: pass diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index b63fd0ecdb1..df219091074 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -4,14 +4,26 @@ Amazon Bedrock Mantle - Responses API backend. gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses` path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides -only the endpoint URL and Bearer authentication. +only the endpoint URL and authentication. -Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the -standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4. +Auth: Bearer token (BEDROCK_MANTLE_API_KEY or the standard +AWS_BEARER_TOKEN_BEDROCK, or litellm_params.api_key) when present; otherwise +AWS SigV4 (service name "bedrock") using the standard credential chain (IAM +role / access key / profile / web identity), signed via the shared +BaseAWSLLM._sign_request after the request body is finalized. """ -from typing import Optional +import re +from typing import Optional, Tuple +from botocore.exceptions import ( + CredentialRetrievalError, + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, +) + +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -29,22 +41,44 @@ _BASE_SUFFIXES_TO_STRIP = ( "/v1", ) +# Standard Mantle host: https://bedrock-mantle..api.aws (group 1 = region). +_MANTLE_HOST_RE = re.compile( + r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE +) + class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): + def __init__(self, aws_signer: Optional[BaseAWSLLM] = None): + super().__init__() + self._aws_signer = aws_signer or BaseAWSLLM() + @property def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + @staticmethod + def _resolve_region(params: dict) -> str: + region = params.get("aws_region_name") + if region: + return region + base = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE") + if base: + match = _MANTLE_HOST_RE.match(base.rstrip("/")) + if match: + return match.group(1) + return ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION_NAME") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + def get_complete_url( self, api_base: Optional[str], litellm_params: dict, ) -> str: - region = ( - get_secret_str("BEDROCK_MANTLE_REGION") - or get_secret_str("AWS_REGION") - or BEDROCK_MANTLE_DEFAULT_REGION - ) + region = self._resolve_region({**litellm_params, "api_base": api_base}) base = ( api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") @@ -55,6 +89,11 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): if base.endswith(suffix): base = base[: -len(suffix)] break + # For the standard Mantle host (including the default-region base that + # responses/main.py auto-injects into litellm_params.api_base), pin to the + # single resolved region so aws_region_name wins; preserve custom proxy hosts. + if _MANTLE_HOST_RE.match(base): + base = f"https://bedrock-mantle.{region}.api.aws" return f"{base}/openai/v1/responses" def validate_environment( @@ -66,12 +105,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") ) - if not api_key: - raise ValueError( - "Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY " - "(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key." - ) - headers["Authorization"] = f"Bearer {api_key}" + if api_key: + headers["Authorization"] = f"Bearer {api_key}" return headers def supports_native_file_search(self) -> bool: @@ -79,3 +114,58 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): def supports_native_websocket(self) -> bool: return False + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + bearer = ( + api_key + or get_secret_str("BEDROCK_MANTLE_API_KEY") + or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + ) + if not bearer: + # SigV4 path. Pin the credential-scope region to the region of the actual + # signing URL (api_base, already region-resolved by get_complete_url) so the + # SigV4 scope and the URL host can never disagree. Resolve from api_base first, + # then fall back to the regular precedence. Also drop any caller Authorization + # so _sign_request's restore-original-Authorization step cannot override the + # SigV4 header. + optional_params = { + **optional_params, + "aws_region_name": self._resolve_region( + {**optional_params, "api_base": api_base} + ), + } + headers = {k: v for k, v in headers.items() if k.lower() != "authorization"} + try: + return self._aws_signer._sign_request( + service_name="bedrock", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=bearer, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + except ( + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, + CredentialRetrievalError, + ) as e: + raise ValueError( + "Bedrock Mantle auth failed: no Bearer token and no usable AWS " + "credentials. Set BEDROCK_MANTLE_API_KEY (or AWS_BEARER_TOKEN_BEDROCK) " + "or pass api_key for Bearer auth, or provide AWS credentials " + "(IAM role / access key / profile / web identity) for SigV4." + ) from e diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 31c772510ba..25424feaeb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2318,6 +2318,31 @@ class BaseLLMHTTPHandler: # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + # Sign after the body is final (post-transform/normalize/extra_body and post + # fake-stream prep) so signed bytes match what we send. No-op for providers + # that inherit the default sign_request. + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2330,22 +2355,14 @@ class BaseLLMHTTPHandler: ) try: - if stream: - # For streaming, use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: return MockResponsesAPIStreamingIterator( @@ -2370,13 +2387,12 @@ class BaseLLMHTTPHandler: call_type=CallTypes.responses.value, ) else: - # For non-streaming requests response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: raise self._handle_error( @@ -2464,6 +2480,28 @@ class BaseLLMHTTPHandler: # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2476,22 +2514,14 @@ class BaseLLMHTTPHandler: ) try: - if stream: - # For streaming, we need to use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: @@ -2518,13 +2548,12 @@ class BaseLLMHTTPHandler: call_type=CallTypes.responses.value, ) else: - # For non-streaming, proceed as before response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: @@ -4005,6 +4034,18 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4018,7 +4059,7 @@ class BaseLLMHTTPHandler: try: response = sync_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: @@ -4088,6 +4129,18 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4101,7 +4154,7 @@ class BaseLLMHTTPHandler: try: response = await async_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 14d950ecdf4..248202b644c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -105,6 +105,16 @@ def _extract_text_from_content(content: object) -> str: return "" +def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Optional[dict[str, Any]]: + merged: dict[str, Any] = {} + present = False + for bag in (request_data.get("metadata"), request_data.get("litellm_metadata")): + if isinstance(bag, Mapping): + present = True + merged.update(bag) + return merged if present else None + + class CrowdStrikeAIDRHandler(CustomGuardrail): """ CrowdStrike AIDR AI Guardrail handler to interact with the CrowdStrike AIDR @@ -312,11 +322,27 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): event_type = "output" hook_name = "apply_guardrail (response)" - ai_guard_payload = { + ai_guard_payload: dict[str, Any] = { "guard_input": guard_input.model_dump(mode="json"), "event_type": event_type, } + model = inputs.get("model") + if model: + ai_guard_payload["model"] = model + + metadata = _merge_metadata_bags(request_data) + if metadata is not None: + user_id = metadata.get("user_api_key_user_id") + if user_id: + ai_guard_payload["user_id"] = user_id + + extra_info: dict[str, str] = {} + user_email = metadata.get("user_api_key_user_email") + if user_email: + extra_info["user_name"] = user_email + ai_guard_payload["extra_info"] = extra_info + ai_guard_response = await self._call_crowdstrike_aidr_guard( ai_guard_payload, hook_name ) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 8473b5e77de..c5715872373 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -518,11 +518,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Check if this is a managed file (base64 encoded unified file ID) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + get_models_from_unified_file_id, ) # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) + target_model_names = ( + get_models_from_unified_file_id(is_managed_file) + if is_managed_file + else [] + ) if is_managed_file and user_api_key_dict is not None: file_content = await self._fetch_managed_file_content( file_id=file_id, @@ -560,6 +566,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): await self._enforce_batch_file_model_access( user_api_key_dict=user_api_key_dict, file_content_as_dict=file_content_as_dict, + target_model_names=target_model_names or None, ) input_file_usage = _get_batch_job_input_file_usage( @@ -595,9 +602,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, user_api_key_dict: UserAPIKeyAuth, file_content_as_dict: List[dict], + target_model_names: Optional[List[str]] = None, ) -> None: - """Reject the batch if the caller is not authorized for every - ``body.model`` named inside the JSONL. + """Reject the batch if the caller is not authorized for the upload target. + + For managed files, ``target_model_names`` (from the unified file id) is + the proxy alias the file was uploaded for and is used directly for auth. + For legacy/non-managed files, falls back to ``body.model`` values in the JSONL. Reuses standard auth helpers so the same model access rules the proxy enforces on `/chat/completions` apply here. @@ -614,9 +625,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): from litellm.proxy.proxy_server import proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache - models = _get_models_from_batch_input_file_content(file_content_as_dict) - if not models: - return + if target_model_names: + models = target_model_names + else: + models = _get_models_from_batch_input_file_content(file_content_as_dict) + if not models: + return team_object = None if ( @@ -647,12 +661,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): llm_model_list = llm_router.model_list if llm_router is not None else None for model in models: - # body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth. model_to_check = model - if llm_router is not None: - proxy_model_name = llm_router.resolve_model_name_from_model_id(model) - if proxy_model_name is not None: - model_to_check = proxy_model_name try: if team_object is not None: try: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a7a0b0f6238..03d0e7ef3ee 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1571,7 +1571,7 @@ class Usage(SafeAttributeModel, CompletionUsage): completion_tokens_details: Optional[ Union[CompletionTokensDetailsWrapper, dict] ] = None, - server_tool_use: Optional[ServerToolUse] = None, + server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, ): @@ -1672,6 +1672,9 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details=_prompt_tokens_details or None, ) + if isinstance(server_tool_use, dict): + server_tool_use = ServerToolUse(**server_tool_use) + if server_tool_use is not None: self.server_tool_use = server_tool_use else: # maintain openai compatibility in usage object if possible @@ -3026,6 +3029,12 @@ class StandardCallbackDynamicParams(TypedDict, total=False): wandb_api_key: Optional[str] weave_project_id: Optional[str] + # Datadog dynamic params + dd_api_key: Optional[str] + dd_site: Optional[str] + dd_agent_host: Optional[str] + dd_agent_port: Optional[str] + # Logging settings turn_off_message_logging: Optional[bool] # when true will not log messages litellm_disabled_callbacks: Optional[List[str]] diff --git a/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py new file mode 100644 index 00000000000..1ee497476d1 --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py @@ -0,0 +1,194 @@ +""" +Tests for team-scoped Datadog callback support. + +Verifies that DataDogLogger can be instantiated with per-team credentials +(dd_api_key, dd_site) instead of relying solely on environment variables, +and that the DataDogHandler correctly resolves and caches per-team loggers. +""" + +from unittest.mock import patch + +import pytest + +from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.integrations.datadog.datadog_team_handler import ( + DataDogHandler, + DatadogLoggingConfig, +) +from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import ( + DynamicLoggingCache, +) +from litellm.types.utils import StandardCallbackDynamicParams + + +@pytest.fixture +def datadog_env(monkeypatch): + """Set global DD env vars for the default/global logger.""" + monkeypatch.setenv("DD_API_KEY", "global_api_key") + monkeypatch.setenv("DD_SITE", "us1.datadoghq.com") + + +class TestDataDogLoggerCredentialKwargs: + """Test that DataDogLogger accepts credentials as kwargs.""" + + def test_init_with_explicit_credentials(self): + """Logger should use explicit kwargs instead of env vars.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_api_key="team_api_key", + dd_site="eu1.datadoghq.com", + ) + + assert logger.DD_API_KEY == "team_api_key" + assert "eu1.datadoghq.com" in logger.intake_url + + def test_init_falls_back_to_env_vars(self, datadog_env): + """Logger should fall back to env vars when no kwargs provided.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + assert logger.DD_API_KEY == "global_api_key" + assert "us1.datadoghq.com" in logger.intake_url + + def test_init_kwargs_override_env_vars(self, datadog_env): + """Explicit kwargs should take precedence over env vars.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_api_key="override_key", + dd_site="ap1.datadoghq.com", + ) + + assert logger.DD_API_KEY == "override_key" + assert "ap1.datadoghq.com" in logger.intake_url + + def test_init_with_agent_credentials(self): + """Logger should use agent mode when dd_agent_host is provided.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_agent_host="dd-agent.local", + dd_agent_port="8125", + dd_api_key="agent_api_key", + ) + + assert "dd-agent.local:8125" in logger.intake_url + assert logger.DD_API_KEY == "agent_api_key" + + def test_init_raises_without_credentials(self, monkeypatch): + """Logger should raise if no credentials are available.""" + monkeypatch.delenv("DD_API_KEY", raising=False) + monkeypatch.delenv("DD_SITE", raising=False) + monkeypatch.delenv("LITELLM_DD_AGENT_HOST", raising=False) + + with pytest.raises(Exception, match="DD_API_KEY"): + with patch("asyncio.create_task"): + DataDogLogger() + + +class TestDataDogHandler: + """Test that DataDogHandler resolves the correct logger per team.""" + + def test_creates_team_logger_with_dynamic_credentials(self, datadog_env): + """Should create a new logger when team credentials are provided.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_api_key="team_a_key", + dd_site="eu1.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result.DD_API_KEY == "team_a_key" + assert "eu1.datadoghq.com" in result.intake_url + + def test_caches_team_logger(self, datadog_env): + """Same team credentials should return the same cached logger instance.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_api_key="team_b_key", + dd_site="us5.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result1 = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + result2 = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result1 is result2 + + def test_different_teams_get_different_loggers(self, datadog_env): + """Different team credentials should create separate logger instances.""" + cache = DynamicLoggingCache() + + params_a = StandardCallbackDynamicParams( + dd_api_key="team_a_key", + dd_site="us1.datadoghq.com", + ) + params_b = StandardCallbackDynamicParams( + dd_api_key="team_b_key", + dd_site="eu1.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result_a = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params_a, + in_memory_dynamic_logger_cache=cache, + ) + result_b = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params_b, + in_memory_dynamic_logger_cache=cache, + ) + + assert result_a is not result_b + assert result_a.DD_API_KEY == "team_a_key" + assert result_b.DD_API_KEY == "team_b_key" + + def test_request_blocked_callback_params_includes_dd(self): + """DD params should be blocked from request-level metadata (security).""" + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + _request_blocked_callback_params, + ) + + assert "dd_api_key" in _request_blocked_callback_params + assert "dd_site" in _request_blocked_callback_params + assert "dd_agent_host" in _request_blocked_callback_params + assert "dd_agent_port" in _request_blocked_callback_params + + +class TestDynamicCredentialDetection: + """Test that _dynamic_datadog_credentials_are_passed works correctly.""" + + def test_no_credentials(self): + params = StandardCallbackDynamicParams() + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is False + + def test_dd_api_key_only(self): + params = StandardCallbackDynamicParams(dd_api_key="key") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + def test_dd_site_only(self): + params = StandardCallbackDynamicParams(dd_site="site") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + def test_dd_agent_host_only(self): + params = StandardCallbackDynamicParams(dd_agent_host="host") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + +class TestStandardCallbackDynamicParamsIncludesDatadog: + """Verify that Datadog params are in the allow-list.""" + + def test_dd_params_in_annotations(self): + annotations = StandardCallbackDynamicParams.__annotations__ + assert "dd_api_key" in annotations + assert "dd_site" in annotations + assert "dd_agent_host" in annotations + assert "dd_agent_port" in annotations diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py new file mode 100644 index 00000000000..4eee6b59d34 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py @@ -0,0 +1,88 @@ +""" +Tests that the cost-tracking call sites tolerate ``server_tool_use`` being +either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, + _get_web_search_requests, +) +from litellm.types.utils import ModelResponse, ServerToolUse, Usage + + +class _UsageWithDictServerToolUse: + """ + Tiny stand-in that mimics the broken streaming-rebuild shape: + ``server_tool_use`` is a plain dict. + """ + + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + self.prompt_tokens_details = None + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 5}) == 5 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + stu = ServerToolUse(web_search_requests=7) + assert _get_web_search_requests(stu) == 7 + + +def test_get_web_search_requests_handles_pydantic_with_none_value(): + stu = ServerToolUse() + assert _get_web_search_requests(stu) is None + + +def test_response_object_includes_web_search_call_with_dict_server_tool_use(): + """ + The exact bug: ``usage.server_tool_use`` is a dict and the check in + ``response_object_includes_web_search_call`` used to crash with + ``AttributeError``. + """ + response = ModelResponse() + usage = _UsageWithDictServerToolUse({"web_search_requests": 2}) + + # Must not raise — and must correctly detect the web search call. + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_pydantic_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2)) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_none_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(None) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is False diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py new file mode 100644 index 00000000000..4e28d5ba7d2 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py @@ -0,0 +1,130 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/26153 + +``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain +``dict`` when reconstructing a streaming response. Downstream cost-calculation +code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call`` +and ``get_cost_for_anthropic_web_search``) accesses +``usage.server_tool_use.web_search_requests`` as an attribute, which raised +``AttributeError: 'dict' object has no attribute 'web_search_requests'``. + +These tests reconstruct streaming chunks for an Anthropic-style web_search +response and assert: + +1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for + ``usage.server_tool_use``. +2. ``completion_cost`` runs end-to-end on the rebuilt response without + raising ``AttributeError``. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm import completion_cost, stream_chunk_builder +from litellm.types.utils import ( + Delta, + ModelResponseStream, + ServerToolUse, + StreamingChoices, + Usage, +) + + +def _make_text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content=text), + ) + ], + ) + + +def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream: + """Final chunk where server_tool_use is a *dict* — reproduces the bug shape.""" + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=Usage( + prompt_tokens=42, + completion_tokens=11, + total_tokens=53, + # NOTE: passed as a dict on purpose — this is the shape that + # historically slipped through stream_chunk_builder unchanged. + server_tool_use={"web_search_requests": 3}, + ), + ) + + +def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic(): + """ + Regression: stream_chunk_builder must produce ServerToolUse, not dict. + """ + chunks = [ + _make_text_chunk("Otters "), + _make_text_chunk("are great."), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + + assert rebuilt is not None + assert rebuilt.usage is not None # type: ignore[attr-defined] + server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined] + + assert ( + server_tool_use is not None + ), "server_tool_use should be carried through from the final chunk" + assert isinstance(server_tool_use, ServerToolUse), ( + f"expected ServerToolUse, got {type(server_tool_use).__name__}: " + f"{server_tool_use!r}" + ) + # Attribute access must not raise (this is exactly what was broken). + assert server_tool_use.web_search_requests == 3 + + +def test_completion_cost_does_not_raise_on_streaming_web_search_response(): + """ + Regression: completion_cost(...) must not raise AttributeError when the + response was reconstructed by stream_chunk_builder from a streaming + Anthropic web_search call. + """ + chunks = [ + _make_text_chunk("hello"), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + assert rebuilt is not None + + # The exact dollar amount depends on the model-pricing table; what matters + # for this regression is that it does NOT raise AttributeError on + # `dict has no attribute 'web_search_requests'`. + try: + cost = completion_cost(completion_response=rebuilt) + except AttributeError as e: # pragma: no cover - regression guard + pytest.fail( + "completion_cost raised AttributeError after stream_chunk_builder " + f"(issue #26153 regression): {e}" + ) + + assert isinstance(cost, (int, float)) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..35aca525f6c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -520,7 +520,10 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use["web_search_requests"] == 2 + # server_tool_use must be a ServerToolUse pydantic so downstream cost-calc + # (which uses attribute access) works. See issue #26153. + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py new file mode 100644 index 00000000000..70fef0162e6 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py @@ -0,0 +1,94 @@ +""" +Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use`` +being either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.anthropic.cost_calculation import ( + _get_web_search_requests, + get_cost_for_anthropic_web_search, +) +from litellm.types.utils import ModelInfo, ServerToolUse + + +class _UsageWithServerToolUse: + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + + +def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo: + info: ModelInfo = { # type: ignore[typeddict-item] + "search_context_cost_per_query": { + "search_context_size_low": cost_per_query, + "search_context_size_medium": cost_per_query, + "search_context_size_high": cost_per_query, + } + } + return info + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 4}) == 4 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2 + + +def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use(): + """ + Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and + direct attribute access on it raised ``AttributeError``. + """ + usage = _UsageWithServerToolUse({"web_search_requests": 3}) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use(): + usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3)) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_none_server_tool_use(): + usage = _UsageWithServerToolUse(None) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == 0.0 + + +def test_get_cost_for_anthropic_web_search_with_no_usage(): + info = _make_model_info(cost_per_query=0.01) + cost = get_cost_for_anthropic_web_search(model_info=info, usage=None) + assert cost == 0.0 diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index e2133d56f89..92b5ca7b10b 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -12,6 +12,11 @@ import sys sys.path.insert(0, os.path.abspath("../../../../..")) import pytest +from botocore.exceptions import ( + ConnectTimeoutError, + PartialCredentialsError, + ProfileNotFound, +) import litellm from litellm.llms.bedrock_mantle.responses.transformation import ( @@ -114,16 +119,15 @@ class TestBedrockMantleResponsesAuth: ) assert headers["Authorization"] == "Bearer bearer-key" - def test_missing_key_raises(self, monkeypatch): + def test_missing_bearer_does_not_raise_in_validate_environment(self, monkeypatch): + # SigV4 may still apply, so validate_environment must defer instead of raising. monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) cfg = BedrockMantleResponsesAPIConfig() - with pytest.raises(ValueError, match="Bedrock Mantle API key"): - cfg.validate_environment( - headers={}, - model="openai.gpt-5.5", - litellm_params=GenericLiteLLMParams(), - ) + headers = cfg.validate_environment( + headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams() + ) + assert "Authorization" not in headers def test_custom_llm_provider(self): cfg = BedrockMantleResponsesAPIConfig() @@ -261,6 +265,386 @@ def local_cost_map(monkeypatch): litellm.get_model_info.cache_clear() +class TestBedrockMantleResponsesSigV4: + def test_bearer_short_circuits_without_credentials(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="bearer-from-config", + ) + assert headers["Authorization"] == "Bearer bearer-from-config" + assert signed_body == b'{"input": "hi"}' + signer.get_credentials.assert_not_called() + + def test_bearer_resolved_from_mantle_env_key(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"] == "Bearer env-bearer" + + def test_bearer_arg_takes_priority_over_mantle_env_key(self, monkeypatch): + # The passed api_key (e.g. litellm_params.api_key) must win over the env + # bearer; a reordered precedence chain would silently use the wrong token. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="arg-bearer", + ) + assert headers["Authorization"] == "Bearer arg-bearer" + signer.get_credentials.assert_not_called() + + def test_access_key_produces_sigv4_headers(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_session_token": "session-token-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Credential=AKIAEXAMPLE/" in headers["Authorization"] + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "X-Amz-Date" in headers + assert headers["X-Amz-Security-Token"] == "session-token-test" + assert signed_body == b'{"input": "hi"}' + + def test_assume_role_path_produces_sigv4_headers(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.credentials import Credentials + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + return_value=Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_session_name": "litellm-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + signer.get_credentials.assert_called_once() + call = signer.get_credentials.call_args.kwargs + assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role" + assert call["aws_session_name"] == "litellm-test" + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + + def test_signed_body_matches_final_data_after_normalize(self, monkeypatch): + """Core regression: the signed bytes must equal the bytes actually sent. + + Sign the *final* data dict and assert the returned signed_body decodes to + exactly that dict, so a later change to the data would break the SigV4 hash. + """ + import json + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + final_data = {"model": "openai.gpt-5.5", "input": "hi", "max_output_tokens": 16} + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + _, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data=final_data, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert signed_body is not None + assert json.loads(signed_body) == final_data + + def test_region_comes_from_optional_params(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "eu-west-1", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses", + api_key=None, + ) + assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"] + + def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch): + """Adversarial-review regression: a caller-supplied aws_region_name (no region + env set) must shape BOTH the URL host and the SigV4 credential scope, or the + request is signed for one region and sent to another -> 401. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params = { + "aws_region_name": "ap-southeast-2", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=None, litellm_params=params) + assert ( + url == "https://bedrock-mantle.ap-southeast-2.api.aws/openai/v1/responses" + ) + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"] + + def test_injected_default_region_base_does_not_override_aws_region_name( + self, monkeypatch + ): + """2nd-round adversarial regression: responses/main.py auto-injects + litellm_params.api_base = https://bedrock-mantle..api.aws/v1 (default + region, ignoring aws_region_name). The config must still pin BOTH the URL host + and the SigV4 scope to aws_region_name, or the IAM deployment 401s. A naive + 'resolve region only when api_base is None' fix would fail this test. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + injected_base = "https://bedrock-mantle.us-east-1.api.aws/v1" # default region + params = { + "aws_region_name": "us-east-2", # what the caller actually wants + "api_base": injected_base, + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=injected_base, litellm_params=params) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "us-east-1" not in headers["Authorization"] + + def test_custom_proxy_host_is_preserved(self, monkeypatch): + """A genuinely custom (non-Mantle) api_base host must be preserved, not rewritten + to a bedrock-mantle host. Only standard Mantle hosts are region-pinned. + """ + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url( + api_base="https://mantle-proxy.internal.example/openai/v1", + litellm_params={"aws_region_name": "us-east-2"}, + ) + assert url == "https://mantle-proxy.internal.example/openai/v1/responses" + + def test_caller_authorization_does_not_override_sigv4(self, monkeypatch): + """Adversarial-review regression: a caller-supplied Authorization header (e.g. + from extra_headers, surviving the relaxed validate_environment) must not clobber + the SigV4 Authorization that _sign_request would otherwise restore. + """ + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={"Authorization": "Bearer stale-caller-token"}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Bearer stale-caller-token" not in headers["Authorization"] + + def test_no_bearer_and_no_credentials_raises_both_paths(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.exceptions import NoCredentialsError + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + @pytest.mark.parametrize( + "cred_error", + [ + PartialCredentialsError(provider="env", cred_var="aws_secret_access_key"), + ProfileNotFound(profile="missing-profile"), + ], + ) + def test_partial_credentials_raises_both_paths(self, monkeypatch, cred_error): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=cred_error) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + def test_sts_transport_error_is_not_masked_as_credentials(self, monkeypatch): + # An AssumeRole / web-identity flow hits STS over the network, so a transient + # connection error must surface as itself, not be rewritten into the + # "no usable AWS credentials" message that would send the user to fix the + # wrong thing. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=ConnectTimeoutError( + endpoint_url="https://sts.us-east-2.amazonaws.com" + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ConnectTimeoutError): + cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + + class TestBedrockMantleResponsesPricing: def test_gpt_5_5_pricing_and_mode(self, local_cost_map): info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5") diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 279e9730e69..7321abcee46 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -742,3 +742,241 @@ async def test_anthropic_post_retry_reserializes_mutated_body(): assert first_sent == prebuilt # attempt 0 used prebuilt assert second_sent == _json.dumps(request_body) # attempt 1 re-serialized assert "MUTATED" in second_sent # ... the mutated body + + +def test_base_responses_config_sign_request_is_noop_by_default(): + """Default responses sign_request must be a no-op: unchanged headers, no signed body. + + Guards the 15 existing responses providers from accidental signing when the + handler starts calling sign_request. + """ + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + cfg = OpenAIResponsesAPIConfig() + headers = {"Authorization": "Bearer sk-existing"} + out_headers, signed_body = cfg.sign_request( + headers=headers, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://api.openai.com/v1/responses", + ) + assert out_headers == {"Authorization": "Bearer sk-existing"} + assert signed_body is None + + +def _make_responses_handler_call(signed_body): + """Drive BaseLLMHTTPHandler.response_api_handler with a fully mocked provider + config + sync client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = {"input": "hi"} + provider_config.should_fake_stream.return_value = False + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + ) + return mock_client.post.call_args.kwargs + + +def test_responses_handler_sends_json_when_not_signed(): + """No-op provider (signed_body is None) -> handler posts json=data, no data= bytes.""" + kwargs = _make_responses_handler_call(signed_body=None) + assert kwargs.get("json") == {"input": "hi"} + assert "data" not in kwargs + + +def test_responses_handler_sends_signed_bytes_when_signed(): + """Signing provider -> handler posts the exact signed bytes via data=, not json=.""" + kwargs = _make_responses_handler_call(signed_body=b'{"input": "hi"}') + assert kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_responses_handler_signs_after_fake_stream_prep_strips_stream(): + """Fake-stream signing-order invariant: the bytes SIGNED must equal the bytes SENT. + + In the streaming + fake-stream path the handler first runs + _prepare_fake_stream_request, which pops "stream" out of the body, and only + then calls sign_request. If signing ran before that pop, the signed body + would still carry "stream" while the body sent over the wire would not, + producing a SigV4 payload-hash mismatch (401) for a real Mantle deployment. + We snapshot request_data at sign time and assert "stream" is already gone. + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = { + "input": "hi", + "stream": True, + } + provider_config.should_fake_stream.return_value = True + provider_config.transform_response_api_response.return_value = ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[], + status="completed", + model="openai.gpt-5.5", + ) + + captured = {} + + def _capture_sign(**kwargs): + captured["request_data"] = dict(kwargs["request_data"]) + return ({"X-Signed": "1"}, b'{"input": "hi"}') + + provider_config.sign_request.side_effect = _capture_sign + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={"stream": True}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + fake_stream=True, + ) + + assert "stream" not in captured["request_data"] + assert "input" in captured["request_data"] + + post_kwargs = mock_client.post.call_args.kwargs + assert post_kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in post_kwargs + assert "stream" in post_kwargs + + +def _make_compact_handler_call(signed_body, is_async): + """Drive (async_)compact_response_api_handler with a fully mocked provider config + + client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle SigV4 / bearer). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + compact_url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses/compact" + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_compact_response_api_request.return_value = ( + compact_url, + {"model": "openai.gpt-5.5", "input": "hi"}, + ) + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + provider_config.transform_compact_response_api_response.return_value = "ok" + + spec = AsyncHTTPHandler if is_async else HTTPHandler + mock_client = MagicMock(spec=spec) + if is_async: + mock_client.post = AsyncMock(return_value=MagicMock()) + else: + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + result = handler.compact_response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=is_async, + ) + if is_async: + asyncio.run(result) + return provider_config, mock_client.post.call_args.kwargs + + +def test_compact_handler_sends_json_when_not_signed(): + """No-op provider on compact (signed_body is None) -> posts json=data, no data= bytes.""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=False + ) + provider_config.sign_request.assert_called_once() + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs + + +def test_compact_handler_sends_signed_bytes_when_signed(): + """Signing provider on compact -> posts the signed bytes via data=, not json=. + + Regression for the adversarial-review finding that /responses/compact bypassed + the SigV4 signing hook, so IAM-only Mantle callers sent unsigned bodies. + """ + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=False + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + # signing must use the compact endpoint as api_base, not the create URL + assert provider_config.sign_request.call_args.kwargs["api_base"].endswith( + "/openai/v1/responses/compact" + ) + + +def test_async_compact_handler_sends_signed_bytes_when_signed(): + """Async compact must sign identically to sync (same omission in the async twin).""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=True + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_async_compact_handler_sends_json_when_not_signed(): + """Async no-op provider on compact -> posts json=data, no data= bytes.""" + _provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=True + ) + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index c58c94cbbc7..607c7b0b32d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -478,3 +478,197 @@ async def test_apply_guardrail_request_skipped_messages_stay_aligned( assert result["texts"][1] == "" assert result["texts"][2] == "Here is my SSN: " assert result["structured_messages"] == inputs["structured_messages"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_user_id_model_and_extra_info( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["model"] == "gpt-4o" + assert payload["extra_info"] == {"user_name": "alice@example.com"} + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_extra_info_when_no_email( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gemini-flash", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gemini-flash", + "litellm_metadata": { + "user_api_key_user_id": "uid-no-email", + "user_api_key_user_email": None, + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-no-email" + assert payload["model"] == "gemini-flash" + assert payload["extra_info"] == {} + + +@pytest.mark.asyncio +async def test_apply_guardrail_no_metadata_skips_user_fields( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + } + request_data = {"messages": inputs["structured_messages"]} + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert "user_id" not in payload + assert "model" not in payload + assert "extra_info" not in payload + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "litellm_metadata, metadata", + [ + ( + None, + { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + ), + ( + {"trace_id": "t1"}, + { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + ), + ( + ["unexpected"], + { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + ), + ( + { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + {"trace_id": "t1"}, + ), + ], + ids=[ + "identity_in_metadata_llm_none", + "identity_in_metadata_llm_user_dict", + "identity_in_metadata_llm_non_mapping", + "identity_in_litellm_metadata", + ], +) +async def test_apply_guardrail_reads_identity_from_either_metadata_bag( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, + litellm_metadata, + metadata, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": litellm_metadata, + "metadata": metadata, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["extra_info"] == {"user_name": "alice@example.com"} diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index ae71d10b378..a6f6e651487 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -713,7 +713,8 @@ async def test_count_input_file_usage_decodes_model_embedded_file_id(): @pytest.mark.asyncio async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(): """After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5). - Auth must check the proxy model_name the key was granted, not the stripped id.""" + Auth must check target_model_names from the unified file id, not reverse-map + the stripped id.""" from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter rate_limiter = _PROXY_BatchRateLimiter( @@ -732,7 +733,6 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( ) mock_router = MagicMock() mock_router.model_list = [] - mock_router.resolve_model_name_from_model_id.return_value = proxy_alias can_key_call_model = AsyncMock(return_value=True) with ( @@ -745,10 +745,105 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, file_content_as_dict=file_dict, + target_model_names=[proxy_alias], ) can_key_call_model.assert_awaited_once() assert can_key_call_model.await_args.kwargs["model"] == proxy_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_list_order", + [ + [ + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + ], + [ + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + ], + [ + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + ], + ], +) +async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup( + model_list_order, +): + """LIT-3593: three deployments strip to gpt-5.5; auth must use the upload + target alias from target_model_names, not first-match reverse lookup.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + batch_alias = "openai/openai/gpt-5.5-batch" + deployment_templates = { + "openai/openai/gpt-5.5": { + "model_name": "openai/openai/gpt-5.5", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + "openai/openai/gpt-5.5-batch": { + "model_name": "openai/openai/gpt-5.5-batch", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"}, + }, + "us/azure/openai/gpt-5.5": { + "model_name": "us/azure/openai/gpt-5.5", + "litellm_params": {"model": "azure/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + } + mock_router = MagicMock() + mock_router.model_list = [deployment_templates[name] for name in model_list_order] + + def _resolve(model_id): + for deployment in mock_router.model_list: + actual_model = deployment.get("litellm_params", {}).get("model") + if actual_model == model_id or ( + actual_model and actual_model.endswith(f"/{model_id}") + ): + return deployment.get("model_name") + return None + + mock_router.resolve_model_name_from_model_id.side_effect = _resolve + + file_dict = [ + {"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}} + ] + user = UserAPIKeyAuth( + api_key="sk-ok", + user_id="alice", + models=[batch_alias], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + can_key_call_model = AsyncMock(return_value=True) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=can_key_call_model, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + ): + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + target_model_names=[batch_alias], + ) + + can_key_call_model.assert_awaited_once() + assert can_key_call_model.await_args.kwargs["model"] == batch_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() @pytest.mark.asyncio