diff --git a/.circleci/config.yml b/.circleci/config.yml index aa851f829e4..32d2cf0390c 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1440,6 +1440,7 @@ jobs: TEST_FILES=$(printf "%s\n" \ tests/local_testing/test_dual_cache.py \ tests/local_testing/test_redis_batch_optimizations.py \ + tests/local_testing/test_redis_increment_with_floor.py \ tests/local_testing/test_router_utils.py) echo "$TEST_FILES" | circleci tests run \ --verbose \ diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 4f56e78ddee..c6901411167 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -117,6 +117,9 @@ jobs: - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + - name: Run pytest tests/test_litellm_rust with the compiled extension + run: make test-rust-extension + - run: >- uv build --wheel --out-dir panic-dist --config-setting "maturin.build-args=--features panic-test,extension-module" diff --git a/Makefile b/Makefile index ab11220821f..91835e19e3c 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,7 @@ .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \ test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \ test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \ + test-rust-extension \ info lint lint-inner lint-dev lint-checks format \ lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ @@ -54,6 +55,7 @@ help: @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" + @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" @echo "" @echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide" @echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine." @@ -289,6 +291,17 @@ pre-commit: @$(MAKE) check # Testing targets +test-rust-extension: + @temporary=$$(mktemp -d) && \ + trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \ + $(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \ + set -- "$$temporary"/wheels/*.whl && \ + [ "$$#" -eq 1 ] && \ + UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \ + $(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \ + LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ + "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust + test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 71e7e9c683b..2145f891318 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -8,14 +8,10 @@ import tempfile import time from dataclasses import dataclass, replace from pathlib import Path -from typing import Optional +from typing import TYPE_CHECKING, Final, Optional from litellm_proxy_extras import prisma_toolchain from litellm_proxy_extras._logging import logger -from litellm_proxy_extras.replica_identity import ( - REPLICA_IDENTITY_FULL_ENV_VAR, - apply_replica_identity_full, -) from litellm_proxy_extras.prisma_toolchain import ( PRISMA_COMMAND_TIMEOUT_ENV_VAR, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, @@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import ( prisma_command_timeout, prisma_migrate_deploy_timeout, ) +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) + +if TYPE_CHECKING: + import psycopg + import psycopg.sql def str_to_bool(value: Optional[str]) -> bool: @@ -46,6 +50,28 @@ def _get_prisma_env() -> dict: _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") _MIGRATION_DEADLOCK_MARKER = "deadlock detected" +INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") +_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") +_INVALID_LITELLM_INDEXES_SQL: Final = ( + "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " + "FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_class t ON t.oid = i.indrelid " + "JOIN pg_namespace n ON n.oid = t.relnamespace " + "WHERE NOT i.indisvalid " + " AND c.relkind = 'i' " + " AND n.nspname = %s " + " AND t.relname LIKE %s " + " AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) " + "ORDER BY c.relname" +) + + +@dataclass(frozen=True, slots=True) +class _InvalidIndex: + schema: str + name: str + table_size: str MAX_MIGRATE_DEPLOY_ATTEMPTS = 4 @@ -624,7 +650,7 @@ class ProxyExtrasDBManager: def _strip_prisma_query_params(url: str) -> str: """Remove Prisma-specific query params (connection_limit, pool_timeout, schema, etc.) from DATABASE_URL so psycopg can parse it.""" - from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode + from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse parsed = urlparse(url) if not parsed.query: @@ -645,7 +671,7 @@ class ProxyExtrasDBManager: "target_session_attrs", } kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params] - return urlunparse(parsed._replace(query=urlencode(kept))) + return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote))) @staticmethod def _warn_if_db_ahead_of_head(migrations_dir: str) -> None: @@ -719,6 +745,95 @@ class ProxyExtrasDBManager: ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""), ) + @staticmethod + def _invalid_litellm_indexes( + conn: "psycopg.Connection[tuple[str, str, str]]", schema: str + ) -> tuple[_InvalidIndex, ...]: + rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall() + return tuple(_InvalidIndex(*row) for row in rows) + + @staticmethod + def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]: + from psycopg import sql + + target: Final = sql.Identifier(index.schema, index.name) + if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name): + return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover" + return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt" + + @staticmethod + def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None: + import psycopg + + statement, action = ProxyExtrasDBManager._index_repair(index) + try: + conn.execute(statement) + except psycopg.Error as e: + logger.warning( + "Could not repair invalid index %s.%s, will retry on the next startup. " + "If this keeps happening, run `%s` by hand as the index owner. Error: %s", + index.schema, + index.name, + statement.as_string(conn), + e, + ) + return + logger.info("%s invalid index %s.%s", action, index.schema, index.name) + + @staticmethod + def repair_invalid_indexes(lock_timeout: str = "30s") -> bool: + """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left + INVALID (a migration deadlock between replicas is the usual cause; the + retried migration skips them because of IF NOT EXISTS). Never raises: + returns True when no invalid index remains, False when the repair was + skipped or failed and will be retried on the next startup. Looks in the + schema DATABASE_URL names, the only URL Prisma migrates through, but + connects over DIRECT_URL when set: the session settings, the advisory + lock and REINDEX CONCURRENTLY all need one server session, which a + transaction pooler does not give.""" + prisma_url: Final = os.getenv("DATABASE_URL") + if not prisma_url: + return False + + try: + import psycopg + from psycopg import sql + except ImportError: + logger.warning( + "psycopg is not installed; skipping the invalid index check. " + "Install the litellm[extra_proxy] extra, which includes psycopg." + ) + return False + + schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) + try: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + conn.execute("SET statement_timeout = 0") + conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) + found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + if not found: + return True + logger.warning( + "Found %d invalid index(es) left by an interrupted CREATE INDEX " + "CONCURRENTLY, rebuilding: %s", + len(found), + ", ".join(f"{index.name} (table size {index.table_size})" for index in found), + ) + lock_row: Final = conn.execute( + "SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,) + ).fetchone() + if lock_row is None or not lock_row[0]: + logger.info("Another replica is already rebuilding the invalid indexes, skipping") + return False + for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema): + ProxyExtrasDBManager._repair_index(conn, index) + remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + except psycopg.Error as e: + logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e) + return False + return not remaining + @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: """ @@ -994,6 +1109,7 @@ class ProxyExtrasDBManager: use_migrate=use_migrate, use_v2_resolver=use_v2_resolver ) if migrated: + ProxyExtrasDBManager.repair_invalid_indexes() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return migrated diff --git a/litellm/batches/main.py b/litellm/batches/main.py index c8360a81c7a..77a4fdebf16 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -319,6 +319,7 @@ def create_batch( timeout=timeout, max_retries=optional_params.max_retries, create_batch_data=_create_batch_request, + custom_endpoint=optional_params.get("custom_endpoint"), ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index aaee7188d86..106c1580110 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -20,6 +20,8 @@ from contextvars import ContextVar from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -80,11 +82,29 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} ) +_INCREMENT_WITH_FLOOR_LUA: Final = ( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end " + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " + "return count" +) + +_LUA_COUNT: Final = TypeAdapter(int) +_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) + + +def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: + return _OPTIONAL_COUNTS.validate_python( + tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) + ) + def _get_call_stack_info(num_frames: int = 2) -> str: """ @@ -736,6 +756,43 @@ class RedisCache(BaseCache): ) raise e + @_redis_circuit_breaker_guard_sync + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call. + + A counter whose key expired while a request was still in flight would otherwise be + recreated negative by that request's decrement. Clamping inside the same call is what + keeps it safe: a separate corrective write could land after another pod's increment and + erase it. + + The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was + created rather than ``ttl`` after it was last touched. Refreshing it on every touch + would keep a count a dead worker never decremented alive for as long as the group + takes traffic. Returns the resulting count. + """ + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval + _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl + ) + return _LUA_COUNT.validate_python(count) + + @_redis_circuit_breaker_guard_sync + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Read integer counters for ``key_list``, in order, raising when Redis cannot answer. + + ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller + cannot tell apart from "every counter is unset". A caller that has to fall back to its + own numbers when Redis is unreachable needs the failure, not a dict of zeros. + """ + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys)) + + @_redis_circuit_breaker_guard + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Async twin of ``batch_get_counts``, raising on failure the same way.""" + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() @@ -1241,6 +1298,14 @@ class RedisCache(BaseCache): result = result.decode() return float(result) + @_redis_circuit_breaker_guard + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl) + return _LUA_COUNT.validate_python(count) + async def flush_cache_buffer(self): print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) diff --git a/litellm/constants.py b/litellm/constants.py index d53686e5e5b..8ef9523a60a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16 MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) +budget_reservation_disabled_info_emitted = False DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) SQS_SEND_MESSAGE_ACTION: Final = "SendMessage" @@ -72,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) +DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float( + os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1") +) DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 16202321709..f9215267bf3 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError): num_retries: int | None = None, headers: dict | None = None, exception_status_code: int | None = None, + response: httpx.Response | None = None, ): request: Final = httpx.Request( method="POST", @@ -352,6 +353,8 @@ class Timeout(openai.APITimeoutError): self.max_retries = max_retries self.num_retries = num_retries self.headers = headers + if response is not None: + self.response = response # custom function to convert to str def __str__(self): diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 3519240dda9..4f9b18713d0 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement @@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient from litellm.litellm_core_utils.prompt_templates.common_utils import ( with_prompt_cache_breakpoint, ) +from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request from litellm.types.integrations.anthropic_cache_control_hook import ( GATEWAY_INJECTED_CACHE_METADATA_KEY, GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT, @@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset( ) OPENAI_API_HOST: Final = "api.openai.com" OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE") +_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues +def _validated_object_mapping(value: object) -> dict[object, object] | None: + try: + return _OBJECT_MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_object_list(value: object) -> list[object] | None: + try: + return _OBJECT_LIST_ADAPTER.validate_python(value) + except ValidationError: + return None + + def supports_openai_prompt_cache_breakpoint(model: str) -> bool: model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model) if model_map_flag is not None: @@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_ class AnthropicCacheControlHook(CustomPromptManagement): + @staticmethod + def _request_value(request_kwargs: object, key: str) -> object: + request_mapping: Final = _validated_object_mapping(request_kwargs) + if request_mapping is None: + return None + return request_mapping.get(key) + + @staticmethod + def _request_user_agent(request_kwargs: object) -> str | None: + proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request") + proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request) + if proxy_server_request_mapping is None: + return None + headers: Final = proxy_server_request_mapping.get("headers") + headers_mapping: Final = _validated_object_mapping(headers) + if headers_mapping is None: + return None + user_agent: Final = next( + (value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"), + None, + ) + return user_agent if isinstance(user_agent, str) else None + + @staticmethod + def _request_system(request_kwargs: object) -> str | list[object] | None: + system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system") + if isinstance(system, str): + return system + return _validated_object_list(system) + def get_chat_completion_prompt( self, model: str, @@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): points: Sequence[CacheControlInjectionPoint], messages: list[AllMessageValues], tools: list[object] | None, + cache_control: object, model: str, custom_llm_provider: str | None, api_base: object, prompt_cache_options: object, ) -> Sequence[Mapping[str, object]] | None: - if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools): + if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control): return None return AnthropicCacheControlHook._stamped_with_dialect( points, model, custom_llm_provider, api_base, prompt_cache_options @@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): messages: list[AllMessageValues], system: str | list | None, tools: list | None, + cache_control: object = None, ) -> bool: """Whether configured injection points must yield to client-set cache_control. @@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement): """ if all(point.get("_litellm_judged") for point in points): return False - return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools) + return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control) @staticmethod def _request_has_cache_control( messages: list[AllMessageValues], system: str | list | None, tools: list | None = None, + cache_control: object = None, ) -> bool: """Return True if the request already carries any client-supplied cache_control. @@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): carry the mark either at the top level (Anthropic shape) or nested under ``function`` (OpenAI shape); the Anthropic chat transform accepts both. """ + if cache_control is not None: + return True if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0: return True if tools is not None: @@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider: str | None, tools: list | None = None, enable_prompt_caching: bool | None = None, + cache_control: object = None, + request_kwargs: object = None, ) -> list[CacheControlInjectionPoint]: """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on. @@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement): if not supports_prompt_caching(model=model, custom_llm_provider=provider): return [] - if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools): + if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control): + return [] + + if is_claude_code_one_shot_subagent_request( + messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs) + ): return [] control: Final = AnthropicCacheControlHook._default_control() @@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): models: Iterable[str], tools: list[AllToolParamValues] | None = None, enable_prompt_caching: bool | None = None, + request_kwargs: object = None, ) -> list[AllMessageValues]: """Return the messages auto prompt caching will send, default breakpoints included. @@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): for candidate in ( AnthropicCacheControlHook.get_default_injection_points( messages=messages, - system=None, model=model, custom_llm_provider=None, tools=tools, enable_prompt_caching=enable_prompt_caching, + system=AnthropicCacheControlHook._request_system(request_kwargs), + cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"), + request_kwargs=request_kwargs, ) for model in models ) @@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): non_default_params["cache_control_injection_points"], messages, tools, + non_default_params.get("cache_control"), model, custom_llm_provider, api_base, @@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider=custom_llm_provider, tools=tools, enable_prompt_caching=enable_prompt_caching, + cache_control=non_default_params.get("cache_control"), + request_kwargs=non_default_params, ) if points: non_default_params["cache_control_injection_points"] = points @@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy bool | None, kwargs.pop("enable_prompt_caching", None) ) + cache_control: Final = kwargs.get("cache_control") configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) ) - if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools): + if configured and AnthropicCacheControlHook._should_stand_down( + configured, typed_messages, system, tools, cache_control + ): return messages, system injection_points: list[CacheControlInjectionPoint] = configured or [] if not injection_points and model is not None: @@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): model=model, custom_llm_provider=custom_llm_provider, enable_prompt_caching=enable_prompt_caching, + cache_control=cache_control, + request_kwargs=kwargs, ) if not injection_points: return messages, system diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2d66a280663..37d6a7e793d 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger): event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None, supported_event_hooks: list[GuardrailEventHooks], ) -> None: + allowed_hooks: Final = frozenset(supported_event_hooks) | ( + frozenset((GuardrailEventHooks.logging_only,)) + if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks + else frozenset() + ) + def _validate_event_hook_list_is_in_supported_event_hooks( event_hook: list[GuardrailEventHooks] | list[str], supported_event_hooks: list[GuardrailEventHooks], @@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger): for hook in event_hook: if isinstance(hook, str): hook = GuardrailEventHooks(hook) - if hook not in supported_event_hooks: + if hook not in allowed_hooks: raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}") if event_hook is None: @@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger): default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default] _validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks) elif isinstance(event_hook, GuardrailEventHooks): - if event_hook not in supported_event_hooks: + if event_hook not in allowed_hooks: raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}") @staticmethod @@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail - def _deployment_pre_call_target(self) -> "CustomLogger": + def _deployment_hook_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self try: @@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - target: Final = self._deployment_pre_call_target() + target: Final = self._deployment_hook_target() if target is not self: kwargs["guardrail_to_apply"] = self result: Final = await target.async_pre_call_hook( @@ -845,7 +851,9 @@ class CustomGuardrail(CustomLogger): return None # CHECK IF GUARDRAIL REJECTS THE REQUEST - result: Final = await self.async_post_call_success_hook( + target: Final = self._deployment_hook_target() + hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data + result: Final = await target.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth( user_id=request_data.get("user_api_key_user_id"), team_id=request_data.get("user_api_key_team_id"), @@ -853,7 +861,7 @@ class CustomGuardrail(CustomLogger): api_key=request_data.get("user_api_key_hash"), request_route=request_data.get("user_api_key_request_route"), ), - data=request_data, + data=hook_request_data, response=response, ) diff --git a/litellm/litellm_core_utils/credential_accessor.py b/litellm/litellm_core_utils/credential_accessor.py index 7071750b970..7b4e8240b69 100644 --- a/litellm/litellm_core_utils/credential_accessor.py +++ b/litellm/litellm_core_utils/credential_accessor.py @@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem class CredentialAccessor: + @staticmethod + def find_credential(credential_name: str) -> CredentialItem | None: + return next( + (credential for credential in litellm.credential_list if credential.credential_name == credential_name), + None, + ) + @staticmethod def get_credential_values(credential_name: str) -> dict: """Safe accessor for credentials.""" - if not litellm.credential_list: - return {} - for credential in litellm.credential_list: - if credential.credential_name == credential_name: - return credential.credential_values.copy() - return {} + credential: Final = CredentialAccessor.find_credential(credential_name) + return {} if credential is None else credential.credential_values.copy() @staticmethod def upsert_credentials(credentials: list[CredentialItem]): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 8f8c955d971..82708d412c9 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -860,6 +860,7 @@ def _map_bedrock_exception( message=mantle_context_window_message, model=model, llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), ) if ( "too many tokens" in error_str @@ -873,6 +874,7 @@ def _map_bedrock_exception( message=f"BedrockException: Context Window Error - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str: raise BadRequestError( @@ -924,12 +926,14 @@ def _map_bedrock_exception( message=f"BedrockException: Timeout Error - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif "Could not process image" in error_str: raise litellm.InternalServerError( message=f"BedrockException - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif hasattr(original_exception, "status_code"): if original_exception.status_code == 500: @@ -937,10 +941,7 @@ def _map_bedrock_exception( message=f"BedrockException - {original_exception.message}", llm_provider="bedrock", model=model, - response=httpx.Response( - status_code=500, - request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), - ), + response=getattr(original_exception, "response", None), ) elif original_exception.status_code == 401: raise AuthenticationError( @@ -969,6 +970,7 @@ def _map_bedrock_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -1001,6 +1003,7 @@ def _map_bedrock_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, exception_status_code=original_exception.status_code, + response=getattr(original_exception, "response", None), ) diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index f1ae6492e4e..d04abcb6e7b 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -1,7 +1,8 @@ +from collections.abc import Mapping from typing import Final -def get_response_headers(_response_headers: dict | None = None) -> dict: +def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict: """ Sets the Appropriate OpenAI headers for the response and forward all headers as llm_provider-{header} @@ -31,7 +32,7 @@ def get_response_headers(_response_headers: dict | None = None) -> dict: return {**llm_provider_headers, **openai_headers} -def _get_llm_provider_headers(response_headers: dict) -> dict: +def _get_llm_provider_headers(response_headers: Mapping[str, str]) -> dict: """ Adds a llm_provider-{header} to all headers that are not already prefixed with llm_provider diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d9424d6a243..2b57883cc13 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -67,6 +67,96 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") +_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" +_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) + + +def is_claude_code_user_agent(user_agent: str) -> bool: + return user_agent.startswith("claude-cli/") + + +def _validated_claude_code_mapping(value: object) -> dict[object, object] | None: + try: + return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_claude_code_list(value: object) -> list[object] | None: + try: + return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None: + stripped: Final = text.strip() + if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX): + return None + fields: Final = tuple( + field + for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";") + if (field := raw_field.strip()) + ) + if not fields or any("=" not in field for field in fields): + return None + parsed_fields: Final = tuple( + (parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),) + ) + if any(not key or not value for key, value in parsed_fields): + return None + return parsed_fields + + +def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None: + if isinstance(system, str): + return (system,) + blocks: Final = _validated_claude_code_list(system) + if blocks is None: + return None + block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks) + if any(block is None for block in block_mappings): + return None + text_values: Final = tuple( + block.get("text") for block in block_mappings if block is not None and block.get("type") == "text" + ) + if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values): + return None + meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip()) + return meaningful_text or None + + +def _is_claude_code_subagent_billing_system(system: object) -> bool: + billing_texts: Final = _claude_code_billing_texts(system) + if billing_texts is None: + return False + billing_fields: Final = tuple( + fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None + ) + if len(billing_fields) != len(billing_texts): + return False + subagent_values: Final = tuple( + value for fields in billing_fields for key, value in fields if key == "cc_is_subagent" + ) + return subagent_values == ("true",) + + +def is_claude_code_one_shot_subagent_request( + messages: list[AllMessageValues], + system: object, + tools: object, + user_agent: str | None, +) -> bool: + only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None + return ( + user_agent is not None + and is_claude_code_user_agent(user_agent) + and not tools + and only_message is not None + and only_message.get("role") == "user" + and _is_claude_code_subagent_billing_system(system) + ) def _strip_bedrock_id_suffixes(model: str) -> str: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py index c1f10c245f8..d9cc65e730f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp( LiteLLM_Proxy_MCP_Handler, ) - mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_references: return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 690040dd93b..6aa17372258 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -667,7 +667,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(response.read())) + raise BedrockError( + status_code=response.status_code, + message=str(response.read()), + headers=response.headers, + response=response, + ) # LOGGING logging_obj.post_call( @@ -690,6 +695,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( status_code=response.status_code, message=f"AgentCore: Failed to read/parse JSON response body: {e}", + headers=response.headers, ) parsed: Final = self._parse_json_response(response_json) @@ -880,7 +886,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(await response.aread())) + raise BedrockError( + status_code=response.status_code, + message=str(await response.aread()), + headers=response.headers, + response=response, + ) # LOGGING logging_obj.post_call( @@ -903,6 +914,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( status_code=response.status_code, message=f"AgentCore: Failed to read/parse JSON response body: {e}", + headers=response.headers, ) parsed: Final = self._parse_json_response(response_json) @@ -1031,6 +1043,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing response: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) def validate_environment( @@ -1046,7 +1059,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return headers def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) def should_fake_stream( self, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index a75124325ae..984ba371898 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -22,7 +22,7 @@ from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token -from ..common_utils import BedrockError, _get_all_bedrock_regions +from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -66,7 +66,12 @@ def make_sync_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(response.read())) + raise BedrockError( + status_code=response.status_code, + message=str(response.read()), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -247,7 +252,12 @@ class BedrockConverseLLM(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -594,7 +604,12 @@ class BedrockConverseLLM(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index e097805f54a..fa24f8be893 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2255,6 +2255,7 @@ class AmazonConverseConfig(BaseConfig): raise BedrockError( message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", status_code=422, + headers=response.headers, ) """ diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index e30ec731d8c..d489e47c3b5 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -470,6 +470,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing response: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) def validate_environment( @@ -485,7 +486,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): return headers def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) def should_fake_stream( self, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index c39c88240c5..5f8a5544d65 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -42,6 +42,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk from ..common_utils import ( BedrockError, build_bedrock_stream_error, + error_response_text, get_bedrock_response_stream_shape, get_bedrock_tool_name, ) @@ -184,7 +185,12 @@ async def make_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=response.text) + raise BedrockError( + status_code=response.status_code, + message=error_response_text(response), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -228,9 +234,16 @@ async def make_call( ) return completion_stream, response.headers + except BedrockError: + raise except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") except Exception as e: @@ -270,7 +283,12 @@ def make_sync_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=response.text) + raise BedrockError( + status_code=response.status_code, + message=error_response_text(response), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -314,9 +332,16 @@ def make_sync_call( ) return completion_stream, response.headers + except BedrockError: + raise except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") except Exception as e: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 04c6ec86a13..5d39b68d9d5 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -247,4 +247,4 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError: """Return the appropriate error class for Bedrock.""" - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index 1671585be2d..4bf1a1cba73 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -182,4 +182,4 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError: """Return the appropriate error class for Bedrock.""" - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index cd8066cda4d..d12c8aee48c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -212,6 +212,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): raise BedrockError( message=f"Error parsing response: {raw_response.text}, error: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) verbose_logger.debug( @@ -241,6 +242,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): raise BedrockError( message=f"Error setting response content: {e}. Response: {completion_response}", status_code=raw_response.status_code, + headers=raw_response.headers, ) # Calculate usage from headers diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 37121d2ece7..a0e32c8aa22 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -295,7 +295,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): try: completion_response: Final = raw_response.json() except Exception: - raise BedrockError(message=raw_response.text, status_code=raw_response.status_code) + raise BedrockError( + message=raw_response.text, + status_code=raw_response.status_code, + headers=raw_response.headers, + ) verbose_logger.debug( "bedrock invoke response % s", json.dumps(completion_response, indent=4, default=str), @@ -363,6 +367,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing={raw_response.text}, Received error={e}", status_code=422, + headers=raw_response.headers, ) try: @@ -384,6 +389,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error parsing received text={outputText}.\nError-{e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) ## CALCULATING USAGE - bedrock returns usage in the headers @@ -431,7 +437,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) @track_llm_api_timing() async def get_async_custom_stream_wrapper( diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py index 311f3a56b84..fb7f2185ec5 100644 --- a/litellm/llms/bedrock/claude_platform/common_utils.py +++ b/litellm/llms/bedrock/claude_platform/common_utils.py @@ -1,7 +1,10 @@ from typing import Final +import httpx + import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.secret_managers.main import get_secret_str CLAUDE_PLATFORM_SERVICE_NAME: Final = "aws-external-anthropic" @@ -15,6 +18,14 @@ def strip_claude_platform_route(model: str) -> str: class BedrockClaudePlatformMixin(BaseAWSLLM): + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + @staticmethod def _get_workspace_id(optional_params: dict, litellm_params: dict) -> str | None: workspace_id = ( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index fe675a30a00..be4f0f32689 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -33,8 +33,53 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues +_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" + + +def error_response_text(response: httpx.Response) -> str: + try: + return response.text + except httpx.ResponseNotRead: + return response.reason_phrase + + +def _synthesize_error_response( + *, status_code: int, headers: dict[str, object] | httpx.Headers, request: httpx.Request | None +) -> tuple[httpx.Request, httpx.Response]: + error_request: Final = request or httpx.Request(method="POST", url=_ERROR_REQUEST_URL) + safe_headers: Final = ( + headers + if isinstance(headers, httpx.Headers) + else tuple((key, value) for key, value in headers.items() if isinstance(value, (str, bytes))) + ) + return error_request, httpx.Response(status_code=status_code, headers=safe_headers, request=error_request) + + class BedrockError(BaseLLMException): - pass + def __init__( + self, + status_code: int, + message: str, + headers: dict[str, object] | httpx.Headers | None = None, + request: httpx.Request | None = None, + response: httpx.Response | None = None, + body: dict[str, object] | None = None, + status_code_is_synthesized: bool = False, + ) -> None: + error_request, error_response = ( + _synthesize_error_response(status_code=status_code, headers=headers, request=request) + if response is None and headers + else (request, response) + ) + super().__init__( + status_code=status_code, + message=message, + headers=headers, + request=error_request, + response=error_response, + body=body, + status_code_is_synthesized=status_code_is_synthesized, + ) _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 2383350b3a3..1fb53f6ff0a 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -102,6 +102,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise BedrockError( status_code=response.status_code, message=error_text, + headers=response.headers, + response=response, ) bedrock_response: Final = response.json() @@ -124,6 +126,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise BedrockError( status_code=e.response.status_code, message=e.response.text, + headers=e.response.headers, + response=e.response, ) except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 5fb86d476f4..d3725434498 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -132,7 +132,12 @@ class BedrockEmbedding(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -161,7 +166,12 @@ class BedrockEmbedding(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index 18d47301ee5..acb0cc8dcb7 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -20,6 +20,7 @@ import httpx from litellm._logging import verbose_logger from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import FileTypes, ImageObject, ImageResponse @@ -228,6 +229,14 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): """ return _supports_nova_canvas_image_edit_from_model_cost(model or "") + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_supported_openai_params(self, model: str) -> list: return [ "n", diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 5c517f2049c..be6489f20ae 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -114,7 +114,12 @@ class BedrockImageEdit(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -156,7 +161,12 @@ class BedrockImageEdit(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index 24e7ba73075..bc9a64f587a 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.stability import ( OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, @@ -84,6 +85,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): return True return False + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_supported_openai_params(self, model: str) -> list: """ Return list of OpenAI params supported by Bedrock Stability. diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index c78e3c147cb..87762b648e0 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -119,7 +119,12 @@ class BedrockImageGeneration(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") ### FORMAT RESPONSE TO OPENAI FORMAT ### @@ -162,7 +167,12 @@ class BedrockImageGeneration(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 6ff9f0155f9..a715d150b4c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( + BedrockError, apply_bedrock_invoke_structured_output, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, @@ -79,6 +80,14 @@ class AmazonAnthropicClaudeMessagesConfig( BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()) + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def __init__(self, **kwargs): BaseAnthropicMessagesConfig.__init__(self, **kwargs) AmazonInvokeConfig.__init__(self, **kwargs) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index d0a3c37ffb3..fb8bc4f191f 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -2,13 +2,14 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, cast +import httpx from httpx import Response from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo +from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo if TYPE_CHECKING: from httpx import URL @@ -18,6 +19,14 @@ if TYPE_CHECKING: class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig): + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in endpoint diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 1f4c81d6491..3b972961940 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -9,12 +9,14 @@ import json import uuid as uuid_lib from typing import Final, cast +import httpx from pydantic import BaseModel from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm from litellm.types.llms.openai import ( OpenAIRealtimeContentPartDone, @@ -121,6 +123,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self._cumulative_usage = BedrockUsageEvent() self._reported_usage = BedrockUsageEvent() + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" return headers diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 4860c99268e..8847381cbc9 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -46,7 +46,12 @@ class BedrockRerankHandler(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -117,7 +122,12 @@ class BedrockRerankHandler(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 920e566c9dd..e7d706c3731 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -39,7 +39,6 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -380,6 +379,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, message=f"AgentCore gateway MCP error: {error}", + headers=raw_response.headers, ) # A failed tools/call is reported in-band, as HTTP 200 with result.isError @@ -389,6 +389,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}", + headers=raw_response.headers, ) text_items: Final = tuple( @@ -440,6 +441,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=502, message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}", + headers=raw_response.headers, ) def get_error_class( @@ -448,7 +450,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): status_code: int, headers: dict, # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict ) -> Exception: - return BaseLLMException( + return BedrockError( status_code=status_code, message=error_message, headers=headers, diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 6940077391f..27c90c9d71e 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBContent, BedrockKBResponse, @@ -38,6 +39,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): BaseVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: return {} diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 64d7ef2bed6..d91157c3d10 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -13,6 +13,8 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the from collections.abc import AsyncIterator, Iterator from typing import Any, Final +import httpx + import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -24,6 +26,8 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams +from ...base_llm.chat.transformation import BaseLLMException +from ...bedrock.common_utils import BedrockError from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import mantle_base_segment @@ -45,6 +49,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): def get_config(cls): return super().get_config() + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers + ) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def _get_openai_compatible_provider_info( self, api_base: str | None, diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 5179c966584..bbbda4d14b6 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -19,11 +19,14 @@ import json from collections.abc import Mapping from typing import Any, Final +import httpx from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock_mantle.common_utils import ( MANTLE_HOST_RE, BedrockMantleAuthMixin, @@ -98,6 +101,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers + ) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_complete_url( self, api_base: str | None, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f561809940..8b6d3a0ceb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + return self._transform_ocr_response( provider_config=provider_config, model=model, @@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + # Use async response transform for async operations return await provider_config.async_transform_ocr_response( model=model, diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 377cd9f3437..a15ea4d845b 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,6 +1,7 @@ import json -from collections.abc import Coroutine +from collections.abc import Coroutine, Sequence from typing import TYPE_CHECKING, Final, Protocol +from urllib.parse import urlparse import httpx from typing_extensions import ReadOnly, TypedDict @@ -12,11 +13,13 @@ from litellm.litellm_core_utils.url_utils import ( safe_get, ) from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler, _get_httpx_client, get_async_httpx_client, ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( VERTEX_CREDENTIALS_TYPES, @@ -55,6 +58,20 @@ class _FetchedResponseView(TypedDict): response: ReadOnly[httpx.Response] +class _VertexEndpointDeployedModel(TypedDict, total=False): + model: ReadOnly[str] + + +class _VertexEndpointResponse(TypedDict, total=False): + deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]] + + +class _VertexEndpointPayloadView(TypedDict): + """Holds one decoded GET endpoints/ response so the payload reads back typed.""" + + payload: ReadOnly[_VertexEndpointResponse] + + def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse: return response.json() @@ -78,7 +95,17 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location: str | None, timeout: float | httpx.Timeout, max_retries: int | None, + custom_endpoint: bool | None = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: + if custom_endpoint: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " + "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " + "use a publisher model or fine-tuned Gemini endpoint deployment instead." + ), + ) sync_handler: Final = _get_httpx_client() access_token, project_id = self._ensure_access_token( @@ -87,6 +114,26 @@ class VertexAIBatchPrediction(VertexLLM): custom_llm_provider="vertex_ai", ) + headers: Final = { + "Content-Type": "application/json; charset=utf-8", + "Authorization": f"Bearer {access_token}", + } + + transformed_batch_request: Final[VertexAIBatchPredictionJob] = ( + VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( + request=create_batch_data, + vertex_project=vertex_project or project_id, + vertex_location=vertex_location or "us-central1", + ) + ) + vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( + vertex_batch_request=transformed_batch_request, + headers=headers, + sync_handler=sync_handler, + api_base=api_base, + vertex_location=vertex_location or "us-central1", + ) + default_api_base: Final = self.create_vertex_batch_url( vertex_location=vertex_location or "us-central1", vertex_project=vertex_project or project_id, @@ -111,17 +158,6 @@ class VertexAIBatchPrediction(VertexLLM): vertex_api_version="v1", ) - headers: Final = { - "Content-Type": "application/json; charset=utf-8", - "Authorization": f"Bearer {access_token}", - } - - vertex_batch_request: Final[VertexAIBatchPredictionJob] = ( - VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( - request=create_batch_data - ) - ) - if _is_async is True: return self._async_create_batch( vertex_batch_request=vertex_batch_request, @@ -142,6 +178,77 @@ class VertexAIBatchPrediction(VertexLLM): ) return vertex_batch_response + @staticmethod + def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str: + """ + Builds the GET url for resolving an endpoint resource (`projects/../endpoints/`). + + A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the + version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a + mount prefix in front of the full default path. The `:operation` suffix convention from + `_check_custom_proxy` does not apply to a plain resource GET. + """ + default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}" + if not api_base: + return default_endpoint_url + api_base_path: Final = urlparse(api_base).path.rstrip("/") + if api_base_path in ("/v1", "/v1beta1"): + return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url) + return api_base.rstrip("/") + urlparse(default_endpoint_url).path + + def _resolve_fine_tuned_endpoint_model( + self, + vertex_batch_request: VertexAIBatchPredictionJob, + headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers + sync_handler: HTTPHandler, + api_base: str | None, + vertex_location: str, + ) -> VertexAIBatchPredictionJob: + """ + A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only + accepts Model resources, so swap the endpoint resource for its deployed tuned model + (`projects/../locations/../models/`) read from GET endpoints/. + """ + model: Final = vertex_batch_request.get("model", "") + if "/endpoints/" not in model: + return vertex_batch_request + + endpoint_url: Final = self._build_endpoint_resolution_url( + api_base=api_base, + model=model, + vertex_location=vertex_location, + ) + # ``api_base`` can come from caller-supplied request kwargs, so wrap the + # fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata + # targets before the bearer token leaves the process (mirrors retrieve_batch). + fetched: Final[_FetchedResponseView] = { + "response": safe_get( + sync_handler, + endpoint_url, + headers=headers, + ) + } + response: Final = fetched["response"] + if response.status_code != 200: + raise VertexAIError( + status_code=response.status_code, + message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}", + ) + + payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()} + deployed_models: Final = payload_view["payload"].get("deployedModels") or () + deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else "" + if not deployed_model: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model " + "resource to run batch predictions against" + ), + ) + resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model} + return resolved_request + async def _async_create_batch( self, vertex_batch_request: VertexAIBatchPredictionJob, diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index f284b47292b..e63c80dd3cf 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -22,6 +22,8 @@ class VertexAIBatchTransformation: def transform_openai_batch_request_to_vertex_ai_batch_request( cls, request: CreateBatchRequest, + vertex_project: str | None = None, + vertex_location: str | None = None, ) -> VertexAIBatchPredictionJob: """ Transforms OpenAI Batch requests to Vertex AI Batch requests @@ -31,7 +33,11 @@ class VertexAIBatchTransformation: if input_file_id is None: raise ValueError("input_file_id is required, but not provided") input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl") - model: Final[str] = cls._get_model_from_gcs_file(input_file_id) + model: Final[str] = cls._get_batch_job_model( + input_file_id=input_file_id, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) output_config: Final[OutputConfig] = OutputConfig( predictionsFormat="jsonl", gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), @@ -188,6 +194,33 @@ class VertexAIBatchTransformation: path_parts: Final = input_file_id.rsplit("/", 1) return path_parts[0] + @classmethod + def _get_batch_job_model( + cls, + input_file_id: str, + vertex_project: str | None, + vertex_location: str | None, + ) -> str: + """ + Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or + the full `projects/../locations/../endpoints/` resource name for a fine-tuned endpoint. + + The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource + to its deployed tuned model (`projects/../locations/../models/`) before sending the job. + """ + parsed_model: Final = cls._get_model_from_gcs_file(input_file_id) + if not parsed_model.startswith("endpoints/"): + return parsed_model + if not vertex_project: + raise VertexAIError( + status_code=400, + message=( + f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require " + "`vertex_project` to build the endpoint resource name" + ), + ) + return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}" + @classmethod def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str: """ @@ -202,6 +235,9 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + Fine-tuned Gemini endpoints are stored as `endpoints/` in the uri and returned + in that form. + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ model: Final = cls._parse_model_from_gcs_file(gcs_file_uri) @@ -210,11 +246,13 @@ class VertexAIBatchTransformation: status_code=400, message=( "Vertex AI batch creation requires the model to be part of `input_file_id`, but " - f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + f"'{gcs_file_uri}' contains no 'publishers//models/' or " + "'endpoints/' path segment. " "Either upload the input file through LiteLLM (POST /v1/files with " "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " "pass a uri of the form " - "gs:////publishers//models//" + "gs:////publishers//models// " + "(or gs:////endpoints// for fine-tuned models)" ), ) return model @@ -222,18 +260,26 @@ class VertexAIBatchTransformation: @classmethod def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: """ - Returns the `publishers//models/` path from a gcs uri, or None if the uri - does not contain one. + Returns the `publishers//models/` or `endpoints/` path from a + gcs uri, or None if the uri does not contain one. + + A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is + used, so a user-configured bucket prefix that happens to contain `endpoints/` cannot + override the model path LiteLLM appended after it. """ - _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") - if not separator: - return None + unquoted_uri: Final = unquote(gcs_file_uri) + _, separator, model_path = unquoted_uri.partition("publishers/") + if separator: + parts: Final = model_path.split("/") + if len(parts) >= 3 and parts[1] == "models" and parts[2]: + return f"publishers/{'/'.join(parts[:3])}" - parts: Final = model_path.split("/") - if len(parts) < 3 or parts[1] != "models" or not parts[2]: - return None + _, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/") + endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else "" + if endpoint_id.isdigit(): + return f"endpoints/{endpoint_id}" - return f"publishers/{'/'.join(parts[:3])}" + return None @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index fe2e0ab6c06..14aebcaabaf 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -370,6 +370,19 @@ def get_vertex_base_model_name(model: str) -> str: return model +def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: + """ + Fine-tuned Gemini deployments are addressed by a numeric endpoint id, + configured as `vertex_ai/` or `vertex_ai/gemini/`. + + Returns the endpoint id, or None when `model` is a regular publisher model. + Mirrors the online chat path in `_get_vertex_url`, which sends numeric + models to `endpoints/{id}` instead of `publishers/google/models/{model}`. + """ + candidate: Final = model.split("/")[-1] if "gemini/" in model else model + return candidate if candidate.isdigit() else None + + def validate_vertex_location(vertex_location: str | None) -> str: """ Validate a Vertex AI location before interpolating it into a request host or diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index b6ad9fbcc04..263956efc9f 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import ( ) from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, + get_vertex_ai_fine_tuned_endpoint_id, ) from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -707,20 +708,39 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, openai_jsonl_content: list[dict[str, Any]], + deployment_model: str | None = None, ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job named as: litellm-vertex-{model}-{uuid} + + The stored model path decides which Vertex model the batch job later executes against, so + `deployment_model` (the deployment's own configured model) wins over the user-supplied + JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no + deployment config. + + Fine-tuned Gemini deployments (numeric endpoint ids) are stored under + `endpoints/` so the batch transformation can round-trip them into a + `projects/../locations/../endpoints/` batch job model instead of a + nonexistent publisher model. """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model") + raw_model: Final = ( + deployment_model.removeprefix("vertex_ai/") + if deployment_model + else openai_jsonl_content[0].get("body", {}).get("model", "") + ) + endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) + model_path: Final = ( + f"endpoints/{endpoint_id}" + if endpoint_id is not None + else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") + ) + safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name - def get_object_name(self, file_data: FileTypes, purpose: str) -> str: + def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str: """ Get the object name for the request. @@ -728,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): upload is never materialized just to derive the GCS object name. """ if purpose == "batch": - ## 1. If jsonl, derive the object name from the first entry's model + ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) if first_entry is not None: - return self._get_gcs_object_name_from_batch_jsonl([first_entry]) + return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model) ## 2. If not jsonl, store under a server-generated managed object name filename, _ = extract_file_metadata(file_data) @@ -761,6 +781,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Get the complete url for the request """ + if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"): + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " + "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " + "remove this deployment from the batch request (e.g. `target_model_names`) or " + "use a publisher model / fine-tuned Gemini endpoint instead." + ), + ) bucket_name = self._get_configured_bucket_name(litellm_params) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) file_data: Final = data.get("file") @@ -769,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - object_name = self.get_object_name(file_data, purpose) + configured_model: Final = litellm_params.get("model") + object_name = self.get_object_name( + file_data, + purpose, + deployment_model=configured_model if isinstance(configured_model, str) else None, + ) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name: Final = encode_gcs_object_name_for_url(object_name) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index ad9622c9cd0..4bac65125b4 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -3108,15 +3108,17 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_agent( user_api_key_auth: UserAPIKeyAuth | None = None, - agent_object_permission=None, + agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, ) -> list[str]: """ Get allowed MCP servers for an agent (from the agent's object_permission). - Returns the MCP servers from the agent's object_permission. - If agent has no object_permission, returns [] (no extra restriction). An entitlement the - agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the - resolver denies. + Returns the agent's direct servers, the servers in its access groups, and the servers reached + through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no + object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that + cannot be read, or a declared toolset that resolves to no grants, raises + ``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the + agent as unrestricted. Args: user_api_key_auth: User auth with agent_id @@ -3126,31 +3128,30 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return [] - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + obj_perm: Final = ( + agent_object_permission + if agent_object_permission is not None + else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + ) if obj_perm is None: return [] try: - direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or [] - if isinstance(direct_mcp_servers, str): - direct_mcp_servers = [] - mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or [] - if isinstance(mcp_access_groups, str): - mcp_access_groups = [] - - # Permission entries may be server_ids OR names/aliases — expand to ids. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers)) - - access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups) - all_servers: Final = expanded_direct_servers + access_group_servers - return list(set(all_servers)) + expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list( + obj_perm.mcp_servers or [] + ) + access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( + obj_perm.mcp_access_groups or [] + ) + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm) + return list({*expanded_direct_servers, *access_group_servers, *toolset_grants}) except Exception as e: + if isinstance(e, UnloadableEntitlementError): + raise verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] @@ -3158,13 +3159,15 @@ class MCPRequestHandler: async def _get_agent_tool_permissions_for_server( server_id: str, user_api_key_auth: UserAPIKeyAuth | None = None, - agent_object_permission=None, + agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, ) -> list[str] | None: """ - Get allowed tool names for a server from the agent's object_permission. - Returns None if agent has no tool restrictions for this server. An entitlement the agent - LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the - tool resolver turns into deny-all for the server rather than an unrestricted tool list. + Get allowed tool names for a server from the agent's object_permission: the union of its + direct tool permissions and the tools its toolsets grant on that server, mirroring the key and + team levels. Returns None if agent has no tool restrictions for this server. An entitlement the + agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises + ``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the + server rather than an unrestricted tool list. Args: server_id: Server ID to check permissions for @@ -3175,24 +3178,30 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return None - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + obj_perm: Final = ( + agent_object_permission + if agent_object_permission is not None + else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + ) if obj_perm is None: return None try: - mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None) - if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict): - return None - # Dict keys may be server_ids OR names/aliases; normalize before lookup. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id) - return list(tools) if tools else None + direct_tools: Final = ( + global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id) + if obj_perm.mcp_tool_permissions + else None + ) + toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id) + agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools) + return list(agent_tools) if agent_tools else None except Exception as e: + if isinstance(e, UnloadableEntitlementError): + raise verbose_logger.warning("Failed to get agent tool permissions for server: %s", e) return None diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 082a90fdcfb..7379126983a 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -4,7 +4,7 @@ import hashlib import json from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -125,6 +125,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): server_id: str +OAuthGrantState = Literal["valid", "refreshable", "absent"] + + class _OAuthTokenRefreshResponse(TypedDict, total=False): access_token: str refresh_token: str @@ -1465,6 +1468,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in return False +def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState: + """Classify local grant readiness without attempting a refresh or checking upstream revocation.""" + if not cred or not cred.get("access_token"): + return "absent" + if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + return "valid" + return "refreshable" if cred.get("refresh_token") else "absent" + + async def get_user_oauth_credential( prisma_client: PrismaClient, user_id: str, @@ -1727,12 +1739,11 @@ async def resolve_valid_user_oauth_token( dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh actually happens, so the valid-token path never requires a DB handle. """ - if not cred or not cred.get("access_token"): + grant: Final = oauth_grant_state(cred) + if cred is None or grant == "absent": return None - if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + if grant == "valid": return cred - if not cred.get("refresh_token"): - return None if prisma_client is None: from litellm.proxy.utils import get_prisma_client_or_throw diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index e10bfd41ed6..cab4b6c161a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + VendorCredentialState, aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -798,21 +800,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302) -async def _bridge_authorize_access_denial( - litellm_user_id: str, - mcp_server: MCPServer, - redirect_uri: str, - state: str, -) -> RedirectResponse | None: - """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed. - - Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the - same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting - session can actually list and call the server's tools. Without this gate the flow completes, the - client shows connected, and every tool request fail-closes to an empty list with nothing telling - the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or - deactivated user denies like a missing grant, fail closed. - """ +async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool: from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) @@ -821,13 +809,22 @@ async def _bridge_authorize_access_denial( ) try: - admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id) + admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id) except HTTPException as exc: if exc.status_code >= 500: raise - return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) - allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted) - if mcp_server.server_id in allowed_server_ids: + return False + return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted) + + +async def _bridge_authorize_access_denial( + litellm_user_id: str, + mcp_server: MCPServer, + redirect_uri: str, + state: str, +) -> RedirectResponse | None: + """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.""" + if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id): return None return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) @@ -1910,6 +1907,38 @@ async def token_endpoint( ) +async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState: + """Whether the gateway itself can see a live vendor credential for this user and server. + + The one reading of "authorized" the connect page displays and the finish step enforces, so + the button a user sees and the grant they get cannot disagree. A read fault is neither, and + fails the scoped grant closed.""" + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # circular import at module load + get_user_oauth_credential, + oauth_grant_state, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load + + if prisma_client is None: + return "unavailable" + try: + credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id) + except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed + return "unavailable" + return "absent" if oauth_grant_state(credential) == "absent" else "present" + + +@router.get("/authorize/flow") +async def authorize_flow(request: Request, flow: str) -> Response: + return await describe_connect_flow( + request=request, + flow_handle=flow, + session_user_id=_session_cookie_user_id(request), + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, + ) + + @router.post("/authorize/complete") async def authorize_complete( request: Request, @@ -1934,6 +1963,8 @@ async def authorize_complete( delivery=delivery, team_id=team_id, decision=decision, + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, ) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index c7b0045dde5..3d94fa345d0 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code" ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"] ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]] -"""Injected live-user revalidation (the token endpoint's mirror of admission): -``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is -a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else -fails the grant closed.""" +VendorCredentialState = Literal["present", "absent", "unavailable"] +"""The per-user vendor credential read has three outcomes: present, absent, or unavailable.""" _DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry" _DB_FAULTED_DESCRIPTION: Final = ( @@ -195,6 +193,16 @@ class ConsentTeam(BaseModel): team_alias: str | None = None +class LookupVendorCredential(Protocol): + """Injected read of a user's vendor credential for one server.""" + + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ... + + +class LookupServerReachability(Protocol): + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ... + + class LookupConsentTeams(Protocol): """Injected lookup of the teams a signed-in user may bind a proxy-API credential to.""" @@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr return "unresolvable" +async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: + return "unavailable" + + +async def _unreachable_server(user_id: str, server_id: str) -> bool: + return False + + class GatewayDcrClient(BaseModel): """The registration record sealed into a gateway DCR ``client_id``. @@ -449,7 +465,10 @@ def aggregate_authorize( A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the flow to that one server: the scope is sealed into the flow, carried into the code, and - bound into the session token, while the connect page interlude runs exactly as before. + bound into the session token. The connect URL carries only the flow handle; the page + learns the client origin, the scoped server, and whether its vendor OAuth is done from + :func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry + steers which server the page authorizes or names on the confirmation. Validation failures respond directly with 400 and never redirect: per RFC 6749 section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and @@ -474,10 +493,7 @@ def aggregate_authorize( resource_server_id=scoped_server.server_id if scoped_server is not None else None, audience=None, ) - connect_url: Final = _append_query_params( - f"{base_url}/ui/connect", - (("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))), - ) + connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) _set_flow_cookie(response, request, handle, flow) return response @@ -684,6 +700,99 @@ def _origin_only(url: str) -> str: return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else "" +def _open_flow_for( + request: Request, flow_handle: str, session_user_id: str | None, now: datetime +) -> _ConnectFlow | Response: + sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) + if sealed_flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) + if flow is None or now.timestamp() >= flow.exp: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + if session_user_id is None: + return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") + if session_user_id != flow.user_id: + return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + return flow + + +async def _flow_target( + flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability +) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]: + if flow.resource_server_id is None: + return "unscoped", None + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle + MCPServerManager, + global_mcp_server_manager, + ) + + server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) + if ( + server is None + or not server.is_gateway_managed_oauth2 + or not await lookup_server_reachability(flow.user_id, server.server_id) + ): + return "stale", None + state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive" + return state, server + + +class ConnectFlowDescription(TypedDict): + """What the connect page is allowed to know about one in-flight flow.""" + + state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]] + client_origin: ReadOnly[str] + server_id: ReadOnly[str | None] + server_name: ReadOnly[str | None] + connected: ReadOnly[bool | None] + + +async def _describe_opened_flow( + flow: _ConnectFlow, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> ConnectFlowDescription | Response: + state, server = await _flow_target(flow, lookup_server_reachability) + if state == "interactive" and server is not None: + credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id) + if credential == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + interactive_description: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": server.server_id, + "server_name": server.server_name or server.alias or server.name, + "connected": credential == "present", + } + return interactive_description + described: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None if server is None else server.server_id, + "server_name": None if server is None else (server.server_name or server.alias or server.name), + "connected": state == "m2m" or None, + } + return described + + +async def describe_connect_flow( + request: Request, + flow_handle: str, + session_user_id: str | None, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> Response: + opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc)) + if isinstance(opened, Response): + return opened + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + return ( + described + if isinstance(described, Response) + else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS) + ) + + async def complete_connect_flow( request: Request, flow_handle: str, @@ -692,56 +801,34 @@ async def complete_connect_flow( delivery: str | None = None, team_id: str | None = None, decision: str | None = None, + lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential, + lookup_server_reachability: LookupServerReachability = _unreachable_server, ) -> Response: - """The deliberate finish step of the connect flow: mint the gateway authorization - code and send the browser back to the client. + """Mint the code only after a deliberate POST by the sealed user. - Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly - per-flow cookie plus an exact match between the signed-in user and the user sealed - into the flow: a link crafted by another party dies here with ``access_denied`` - instead of minting a code for the victim's identity. The flow is single-use (an atomic - claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in. - - ``delivery`` chooses how the code reaches the client. Default (absent or - ``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"`` - renders the callback URL on a page instead, for a client whose redirect URI is a - loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box, - container): the 303 would dereference the browser machine's loopback and the code - would never arrive, so the user carries it over by pasting the URL into the client or - fetching it from the client machine's terminal. Manual delivery is honored only for - loopback redirect URIs; a routable redirect URI works from any browser by - construction, so those flows always redirect. The user who sees the page is exactly - the user the 303 would have carried the code to, and the same user already sees the - code today in the dead redirect's address bar, so the page exposes the code to no new - party. Unknown ``delivery`` values are rejected rather than defaulted: a client that - asked for manual delivery and got a dead redirect instead would silently lose its - code. - - ``decision`` and ``team_id`` come from the native-client consent page. ``"deny"`` - burns the flow and sends the client ``error=access_denied`` so it stops waiting; - ``team_id`` is sealed into the code only for proxy-API flows, where it picks which of - the user's teams the minted credential is attributed to. + A scoped flow additionally requires its sealed server to have a live vendor credential + before a code can be minted. The check happens before the single-use claim, so a + premature submit can be retried after authorization; denial deliberately bypasses it. """ if delivery not in (None, "redirect", "manual"): return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'") if decision not in (None, "approve", "deny"): return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'") - sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) - if sealed_flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") - flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) - if flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") now: Final = datetime.now(timezone.utc) - if now.timestamp() >= flow.exp: - return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection") - if session_user_id is None: - return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") - if session_user_id != flow.user_id: - return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + opened: Final = _open_flow_for(request, flow_handle, session_user_id, now) + if isinstance(opened, Response): + return opened + if decision != "deny": + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + if isinstance(described, Response): + return described + if described["state"] == "stale": + return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available") + if described["connected"] is False: + return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") flow_refusal: Final = _claim_refusal( await _SingleUseGuard(cache).claim( - f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS ), replayed=_oauth_error( 400, "invalid_request", "this connect flow was already completed; restart the connection" @@ -750,7 +837,7 @@ async def complete_connect_flow( if flow_refusal is not None: return flow_refusal response: Final = ( - _denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now) + _denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now) ) path, secure = _cookie_path_and_secure(request) response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax") diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0b424c31c4b..a44a041fdff 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -23,6 +23,7 @@ from pydantic import AnyUrl, ConfigDict from starlette.requests import Request as StarletteRequest from starlette.responses import JSONResponse from starlette.types import Message, Receive, Scope, Send +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG @@ -816,6 +817,11 @@ if MCP_AVAILABLE: } } return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) + except HTTPException as e: + from mcp.shared.exceptions import McpError + from mcp.types import INVALID_REQUEST, ErrorData + + raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e except Exception as e: verbose_logger.exception("Error in list_tools endpoint: %s", e) # Return empty list instead of failing completely @@ -1095,6 +1101,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + client_ip=_client_ip, host_progress_callback=host_progress_callback, **data, # for logging ) @@ -1128,7 +1135,7 @@ if MCP_AVAILABLE: except HTTPException as e: verbose_logger.error("HTTPException in MCP tool call: %s", e) return CallToolResult( - content=[TextContent(text=f"Error: {e.detail}", type="text")], + content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")], isError=True, ) except MCPUpstreamAuthError as e: @@ -1392,7 +1399,7 @@ if MCP_AVAILABLE: ######################################################## async def _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers: list[str] | None, + mcp_servers: Sequence[str] | None, allowed_mcp_servers: list[MCPServer], ) -> list[MCPServer]: """ @@ -1413,13 +1420,10 @@ if MCP_AVAILABLE: server_name_matched = False for server in allowed_mcp_servers: - if server: - match_list = [s.lower() for s in iter_known_server_prefixes(server) if s] - - if server_or_group.lower() in match_list: - filtered_server[server.server_id] = server - server_name_matched = True - break + if server and _server_answers_to(server, server_or_group): + filtered_server[server.server_id] = server + server_name_matched = True + break if not server_name_matched: try: @@ -1449,6 +1453,72 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def _http_detail_message(detail: object) -> str: + return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail) + + def _server_answers_to(server: MCPServer, name: str) -> bool: + requested: Final = name.lower() + return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known) + + class _McpDeniedDetail(TypedDict): + error: ReadOnly[str] + + async def raise_denied_scoped_mcp_access( + requested_names: Sequence[str], + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None = None, + ) -> None: + """A scoped request (``/mcp/`` path or ``x-mcp-servers`` header) resolved to zero + allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy + server with no tools. Unknown, unauthorized, and access-group names all share one generic + error so scoping cannot probe which servers exist; the agent variant fires only when the + same request resolves once the agent binding is stripped, proving the binding caused the veto.""" + agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None + if user_api_key_auth is not None and agent_id: + resolved_without_agent: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})), + mcp_servers=requested_names, + client_ip=client_ip, + ) + + def _resolved_to_server(name: str) -> bool: + return any(_server_answers_to(server, name) for server in resolved_without_agent) + + vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None) + if vetoed_server is not None: + agent_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP server '{vetoed_server}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include this server. Add the server " + f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=agent_denial) + vetoed_group: Final = next( + ( + name + for name in requested_names + if not _resolved_to_server(name) + and any(name in (server.access_groups or ()) for server in resolved_without_agent) + ), + None, + ) + if vetoed_group is not None: + group_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the " + f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=group_denial) + generic_denial: Final[_McpDeniedDetail] = { + "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}" + } + raise HTTPException(status_code=403, detail=generic_denial) + def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool: """ Check if a tool name matches any name in the filter list. @@ -1541,7 +1611,7 @@ if MCP_AVAILABLE: async def _get_allowed_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, - mcp_servers: list[str] | None, + mcp_servers: Sequence[str] | None, client_ip: str | None = None, ) -> list[MCPServer]: """Return allowed MCP servers for a request after applying filters. @@ -1977,6 +2047,12 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) # Pre-fetch OAuth credentials only when at least one server uses OAuth2, # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. @@ -2404,6 +2480,8 @@ if MCP_AVAILABLE: ) verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools)) return listing + except HTTPException: + raise except Exception as e: verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) # Continue with an empty listing instead of failing completely @@ -3086,6 +3164,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, **kwargs: Any, ) -> CallToolResult: """ @@ -3116,6 +3195,12 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, allowed_mcp_servers=allowed_mcp_servers, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) if not allowed_mcp_servers: raise HTTPException( status_code=403, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index f19340d30cb..24e2f5ce64d 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -366,6 +366,7 @@ async def handle_mcp_tool_call( from litellm.proxy._experimental.mcp_server.server import ( _get_allowed_mcp_servers, execute_mcp_tool, + raise_denied_scoped_mcp_access, ) allowed_mcp_servers: Final = await _get_allowed_mcp_servers( @@ -373,6 +374,12 @@ async def handle_mcp_tool_call( mcp_servers=mcp_servers, client_ip=client_ip, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_dict, + client_ip=client_ip, + ) # Reject before dispatch when the key has no accessible servers; otherwise an # unprefixed local tool name would fall through to the local registry in diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c71761a7adb..71475320c2c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2634,6 +2634,20 @@ ], "title": "Mcp Tool Permissions" }, + "mcp_toolsets": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Mcp Toolsets" + }, "models": { "anyOf": [ { @@ -19872,6 +19886,46 @@ ] } }, + "/authorize/flow": { + "get": { + "operationId": "authorize_flow_authorize_flow_get", + "parameters": [ + { + "in": "query", + "name": "flow", + "required": true, + "schema": { + "title": "Flow", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Authorize Flow", + "tags": [ + "mcp_discoverable" + ] + } + }, "/callback": { "get": { "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index abce10690e5..4dbae6394f6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2833,7 +2833,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "Enable only if your deployment is experiencing phantom " "BudgetExceededError responses caused by leaked reservations " "(see GitHub issue #27639). " - "A proxy-level WARNING is logged on every request while this flag " + "An INFO notice is logged once per worker at config load while this flag " "is active as a reminder that hard enforcement is relaxed." ), ) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 1e4836654a1..f78c4221f5a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -11,7 +11,7 @@ from fastapi import HTTPException, Request, status from pydantic import PositiveInt, TypeAdapter, ValidationError import litellm -from litellm import Router, provider_list +from litellm import Router, constants, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import ( BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, @@ -1390,6 +1390,24 @@ def warn_once_if_custom_auth_skips_common_checks( _custom_auth_common_checks_warning_emitted = True +def log_once_if_budget_reservation_disabled( + *, + disabled: bool, + logger: Logger = verbose_proxy_logger, +) -> None: + if constants.budget_reservation_disabled_info_emitted or not disabled: + return + logger.info( + "disable_budget_reservation is enabled: skipping optimistic budget " + "reservation. Budget enforcement is read-time only. Concurrent " + "requests can each pass the spend check before their cost is recorded, " + "so a configured budget may be briefly exceeded under high concurrency. " + "Set disable_budget_reservation to False or remove it to restore " + "hard per-request budget enforcement." + ) + constants.budget_reservation_disabled_info_emitted = True # rebind-ok: process-wide one-shot sentinel + + def is_pass_through_provider_route(route: str) -> bool: PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES: Final = [ "vertex-ai", diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 93293db24c6..b39b1f330b3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2706,14 +2706,6 @@ async def _reserve_budget_after_common_checks( if skip_budget_checks: return if general_settings.get("disable_budget_reservation") is True: - verbose_proxy_logger.warning( - "disable_budget_reservation is enabled: skipping optimistic budget " - "reservation. Budget enforcement is read-time only — concurrent " - "requests can each pass the spend check before their cost is recorded, " - "so a configured budget may be briefly exceeded under high concurrency. " - "Set disable_budget_reservation to False or remove it to restore " - "hard per-request budget enforcement." - ) return from litellm.proxy.spend_tracking.budget_reservation import ( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9720e4b1cf8..d0e3914e9fb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3508,9 +3508,13 @@ class ProxyBaseLLMRequestProcessing: error_body: Final = await http_status_error.response.aread() error_text: Final = error_body.decode("utf-8") + error_headers: Final = { # mutable-ok: HTTPException takes a plain header dict + k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items() + } raise HTTPException( status_code=http_status_error.response.status_code, detail={"error": error_text}, + headers=error_headers, ) error_msg: Final = f"{e}" # Check for AttributeError in the exception chain. diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 460b348e188..63da5d15207 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -125,6 +125,7 @@ async def _resync_model_deployments(model_name: str) -> bool: ) return proxy_server.llm_router is not None async with proxy_server.MODEL_RECONCILE_LOCK: + await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client) proxy_server.proxy_config._add_deployment(db_models=rows) proxy_server.llm_model_list = router.get_model_list() return True diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index b7a07d4eeea..6e2a06c96e1 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]: def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: """Checks for differences between current database and Prisma schema. + + Never raises: a diff that cannot be produced, because the runner is missing, + because the command failed, or because it outlived its budget, is reported as + "no diff" so boot continues. + Returns: A tuple containing: - A boolean indicating if differences were found (True) or not (False). - - A string with the diff output or error message. - Raises: - subprocess.CalledProcessError: If the Prisma command fails. - Exception: For any other errors during execution. + - The SQL commands that would close the diff, empty when there is none. """ - verbose_logger.debug("Checking for Prisma schema diff...") try: - result: Final = subprocess.run( + from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" + ) + return False, [] + + verbose_logger.debug("Checking for Prisma schema diff...") + timeout: Final = prisma_command_timeout() + try: + result: Final = run_prisma( [ "prisma", "migrate", @@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: "./schema.prisma", "--script", ], - capture_output=True, - text=True, - check=True, + timeout=timeout, + env=os.environ.copy(), ) - # return True, "Migration diff generated successfully." sql_commands: Final = extract_sql_commands(result.stdout) if sql_commands: @@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: return True, sql_commands else: return False, [] + except subprocess.TimeoutExpired: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Timed out after {timeout}s generating the migration diff. " + f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer." + ) + return False, [] except subprocess.CalledProcessError as e: error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" print(error_message) # noqa: T201 diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 2190ae55fd2..9ea2432f2f5 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -937,9 +937,17 @@ class PrismaManager: use_v2_resolver=use_v2_resolver, ) else: + try: + from litellm_proxy_extras.prisma_toolchain import ( + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e) + return False + PrismaManager._raise_if_partitioned_spend_logs() - # Use prisma db push with increased timeout - subprocess.run( + run_prisma( [ "prisma", "db", @@ -947,13 +955,15 @@ class PrismaManager: "--accept-data-loss", "--skip-generate", ], - timeout=60, - check=True, + timeout=prisma_command_timeout(), + env=os.environ.copy(), + stdout=None, + stderr=None, ) PrismaManager._apply_replica_identity_full_if_requested() return True - except subprocess.TimeoutExpired: - verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1) + except subprocess.TimeoutExpired as e: + verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index a8b33109900..e20f0b320b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -547,7 +547,7 @@ class ToolPermissionGuardrail(CustomGuardrail): for _tool_call, is_allowed, _rule_id, message in checked: if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=message, blocked_content=True @@ -809,7 +809,7 @@ class ToolPermissionGuardrail(CustomGuardrail): new_tools: Final = self._collect_request_tools(data) if not new_tools: - verbose_proxy_logger.warning( + verbose_proxy_logger.debug( "Tool Permission Guardrail: not running guardrail. No tools or functions in data" ) return data @@ -820,7 +820,7 @@ class ToolPermissionGuardrail(CustomGuardrail): is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise HTTPException( status_code=400, diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 8b82842353c..bec12ba8201 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException @@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Parse to separate MCP tools from other tools - mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools: return [] @@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger): return [name for name in names if name] @staticmethod - def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]: + async def _narrow_mcp_references( + tools: Sequence[Mapping[str, object]], + selected_tool_names: list[str], + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None, + ) -> list[object]: """ Restrict each litellm_proxy MCP reference to the semantically selected tools. @@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger): LiteLLM_Proxy_MCP_Handler, ) + via_gateway: Final = await ( + LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools, served_names) + if served_names is not None + else LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools) + ) return [ - ( - {**tool, "allowed_tools": selected_tool_names} - if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool]) - else tool - ) - for tool in tools + {**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool + for tool, routed in zip(tools, via_gateway, strict=True) ] def _is_mcp_tool(self, tool: object) -> bool: @@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger): filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools) selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools) - narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names) + narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names) data["tools"] = narrowed_tools self._emit_filter_metadata_safe( data=data, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e4bce4378f0..0534a792fee 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -789,12 +789,6 @@ def apply_missing_session_id_policy( ) -def is_claude_code_user_agent(user_agent: str) -> bool: - """Claude Code identifies itself as ``claude-cli/ ...``; the IDE - extensions and the Agent SDK run through the same CLI and share that prefix.""" - return user_agent.startswith("claude-cli/") - - def is_codex_user_agent(user_agent: str) -> bool: """Codex builds its user agent as ``/ ...`` and ships several first-party originators: ``codex-tui``, ``codex_cli_rs``, @@ -811,6 +805,8 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c requests routed to providers that reject them. An explicit drop_params from the caller or in the operator's ``litellm_settings`` always wins over this default.""" + from litellm.llms.anthropic.common_utils import is_claude_code_user_agent + if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)): return False if "drop_params" in data: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b95547e2b54..2ea46b740a8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -946,7 +946,14 @@ async def handle_bedrock_count_tokens( except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) - raise HTTPException(status_code=e.status_code, detail={"error": e.message}) + from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers + + provider_headers: Final = getattr(getattr(e, "response", None), "headers", None) + raise HTTPException( + status_code=e.status_code, + detail={"error": e.message}, + headers=get_response_headers(provider_headers) if provider_headers else None, + ) except HTTPException: # Re-raise HTTP exceptions as-is raise diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9bc10949e9f..7bcb79cefc9 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -20,7 +20,11 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import independent_snapshot +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, + independent_snapshot, +) from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, @@ -30,7 +34,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStep, PipelineStepResult, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import GenericGuardrailAPIInputs, StandardLoggingGuardrailInformation if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -258,6 +262,7 @@ class PipelineExecutor: return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="block", step_results=step_results, @@ -267,6 +272,7 @@ class PipelineExecutor: ) if action == "modify_response": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="modify_response", step_results=step_results, @@ -357,16 +363,17 @@ class PipelineExecutor: verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) return ("error", None, f"Guardrail '{step.guardrail}' not found", None) + hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot) + snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input)) + + # Use unified_guardrail path if callback implements apply_guardrail + target: CustomLogger = callback + use_unified: Final = PipelineExecutor.supports_unified_execution(callback) + if use_unified and streaming_chunks is None: + hook_input["guardrail_to_apply"] = callback + target = UnifiedLLMGuardrails() + try: - hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot) - - # Use unified_guardrail path if callback implements apply_guardrail - target: CustomLogger = callback - use_unified: Final = PipelineExecutor.supports_unified_execution(callback) - if use_unified and streaming_chunks is None: - hook_input["guardrail_to_apply"] = callback - target = UnifiedLLMGuardrails() - if mode == "pre_call": response = await target.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -430,6 +437,12 @@ class PipelineExecutor: else: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + finally: + if hook_input is not data: + _append_guardrail_information( + request_data=data, + entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:], + ) @staticmethod def supports_unified_execution(callback: CustomGuardrail) -> bool: @@ -486,6 +499,40 @@ def _restore_request_guardrails( return {**working_data, "metadata": stripped} # mutable-ok: request dict +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]: + bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source)) + recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None + return recorded if isinstance(recorded, list) else [] + + +def _append_guardrail_information( + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data + entries: Sequence[StandardLoggingGuardrailInformation], +) -> None: + if not entries: + return + _, request_bucket = get_or_create_metadata_bucket(request_data) + existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY) + if isinstance(existing, list): + existing.extend(entries) + return + request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries) + + +def _carry_working_guardrail_information( + working_data: Mapping[str, object], + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data +) -> None: + recorded: Final = _recorded_guardrail_information(working_data) + existing: Final = _recorded_guardrail_information(request_data) + if recorded is existing: + return + _append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing]) + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0915b8dd1b9..1d6baf5fded 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -306,6 +306,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.auth_utils import ( check_response_size_is_safe, is_request_body_safe, + log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check @@ -5653,6 +5654,10 @@ class ProxyConfig: run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)), ) + log_once_if_budget_reservation_disabled( + disabled=general_settings.get("disable_budget_reservation") is True, + ) + custom_key_generate: Final = general_settings.get("custom_key_generate", None) if custom_key_generate is not None: user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path) @@ -7104,11 +7109,10 @@ class ProxyConfig: ], ) - # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) - if self._should_load_db_object(object_type="models"): - new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) - - # update llm router + load_models: Final = self._should_load_db_object(object_type="models") + new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None + await self.get_credentials(prisma_client=prisma_client) + if load_models: still_desired_ids = await self._update_llm_router( new_models=new_models, proxy_logging_obj=proxy_logging_obj ) @@ -7148,12 +7152,9 @@ class ProxyConfig: async def _resync_config_from_db() -> None: await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) - async def _resync_credentials_from_db() -> None: - await self.get_credentials(prisma_client=prisma_client) - subscriber: Final = ConfigSyncSubscriber( redis_cache=redis_cache, - resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), + resync_callbacks=(_resync_config_from_db,), ) self.config_sync_subscriber = subscriber subscriber.start() @@ -8008,7 +8009,7 @@ class ProxyConfig: async def get_credentials(self, prisma_client: PrismaClient): try: - credentials = await CredentialsRepository(prisma_client).find_all() + credentials = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_all() credentials = [self.decrypt_credentials(cred) for cred in credentials] await self.delete_credentials(credentials) # delete credentials that are not in the all-up list CredentialAccessor.upsert_credentials(credentials) # upsert credentials that are in the all-up list @@ -9592,19 +9593,6 @@ class ProxyStartupEvent: ) if store_model_in_db is True: - ### GET STORED CREDENTIALS ### - scheduler.add_job( - proxy_config.get_credentials, - "interval", - seconds=config_reload_interval_seconds, - # REMOVED jitter parameter - major cause of memory leak - args=[prisma_client], - id="get_credentials_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - await proxy_config.get_credentials(prisma_client=prisma_client) - # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum # Frequent polling was causing excessive memory allocations scheduler.add_job( @@ -9618,7 +9606,7 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - # this will load all existing models on proxy startup + # this will load all existing credentials and models on proxy startup await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) proxy_config.start_config_sync_subscriber( diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index c541b9b40e5..950fcca2039 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -299,6 +299,8 @@ def compute_autorouter_savings( selected_info: ModelInfo | None = None, baseline_info: ModelInfo | None = None, cost_breakdown: Mapping[str, object] | None = None, + baseline_deployment_id: str | None = None, + selected_deployment_id: str | None = None, ) -> float: """Net dollars the router saved, or cost, by serving this request on ``selected_model``. @@ -334,11 +336,12 @@ def compute_autorouter_savings( selected: Final = _resolve_model(selected_model, selected_provider) if baseline is None or selected is None: return 0.0 - # Same model is only the same cost when it is also the same deployment. Two - # deployments of one model can carry different negotiated rates, and routing from - # the dear one to the cheap one is a real saving that short-circuiting on the model - # name alone reports as zero. - if baseline == selected: + same_target: Final = ( + baseline_deployment_id == selected_deployment_id + if baseline_deployment_id and selected_deployment_id + else baseline == selected + ) + if same_target: return 0.0 basis: Final = _pricing_basis(cost_breakdown) effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline) @@ -517,6 +520,8 @@ def autorouter_savings_for_request( selected_info=_effective_model_info(router_instance, model_id, model or ""), baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""), cost_breakdown=cost_breakdown, + baseline_deployment_id=baseline_id, + selected_deployment_id=model_id, ) classifier_cost: Final = classifier_cost_from_decision(decision) return gross if classifier_cost is None else gross - classifier_cost diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5e74b7324b4..fa14eb5f3c4 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -177,7 +177,7 @@ async def aresponses_api_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform) # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata) @@ -236,6 +236,7 @@ async def aresponses_api_with_mcp( "timeout": timeout, "custom_llm_provider": custom_llm_provider, **kwargs, + "_skip_mcp_handler": True, } # Handle MCP streaming if requested @@ -898,13 +899,14 @@ def _responses_try_dispatch_mcp_gateway( custom_llm_provider: str | None, kwargs: dict[str, object], _is_async: bool, + skip_mcp_handler: bool, ) -> Any | None: """Return a response when MCP gateway handles the call; otherwise None.""" from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): + if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): return None mcp_call_kwargs: Final = { "input": input, @@ -1074,6 +1076,7 @@ def responses( litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) _is_async: Final = kwargs.pop("aresponses", False) is True + skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs) client_headers: Final = kwargs.get("headers") @@ -1168,6 +1171,7 @@ def responses( custom_llm_provider=custom_llm_provider, kwargs=kwargs, _is_async=_is_async, + skip_mcp_handler=skip_mcp_handler, ) if _mcp_dispatch is not None: return _mcp_dispatch diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index a75b3768636..ae18d5f6f1b 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -106,7 +106,7 @@ async def acompletion_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools_with_litellm_proxy: # No MCP tools, proceed with regular completion @@ -114,6 +114,7 @@ async def acompletion_with_mcp( model=model, messages=messages, tools=tools, + _skip_mcp_handler=True, **kwargs, ) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 15434bedbb7..a5021e2f777 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,6 +1,6 @@ import re import traceback -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.utils import ( + iter_known_server_prefixes, logging_safe_mcp_headers, split_server_prefix_from_name, strip_known_server_prefix, @@ -23,6 +24,7 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamingResponse, ) from litellm.types.llms.openai import ToolParam as ResponsesToolParam +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ( CallTypes, ChatCompletionMessageCustomToolCall, @@ -45,6 +47,7 @@ else: # NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling ToolParam: TypeAlias = Mapping[str, object] +SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]] class MCPToolResult(TypedDict): @@ -56,14 +59,74 @@ class MCPToolResult(TypedDict): LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" -# Matches any URL whose path ends with /mcp/ — covers both root-path -# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments. -# A false-positive match (e.g. an external URL that happens to end with /mcp/) results -# in a "server not found" error from the internal gateway, not a silent failure or data leak, -# so this broad pattern is intentional and preferred over anchoring to localhost only. _PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") +def _mcp_server_url(tool: ToolParam) -> str | None: + if not isinstance(tool, dict) or tool.get("type") != "mcp": + return None + server_url: Final = tool.get("server_url") + return server_url if isinstance(server_url, str) else None + + +def _names_gateway_explicitly(tool: ToolParam) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def _proxy_path_mcp_name(tool: ToolParam) -> str | None: + server_url: Final = _mcp_server_url(tool) + match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) + return None if match is None else match.group(1) + + +def _registered_mcp_servers() -> Collection[MCPServer]: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return global_mcp_server_manager.get_registry().values() + + +def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: + requested: Final = name.lower() + return any( + requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name)) + or name in (server.access_groups or ()) + for server in servers + ) + + +async def _toolset_exists(name: str) -> bool: + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return False + return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None + except Exception as e: + verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e) + return False + + +async def _gateway_served_names( + names: Collection[str], + servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, + toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists, +) -> frozenset[str]: + registered: Final = tuple(servers()) if names else () + return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) + + +async def _served_mcp_path_names( + tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] +) -> frozenset[str]: + names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + return await served_names(names) if names else frozenset[str]() + + class LiteLLM_Proxy_MCP_Handler: """ Helper class with static methods for MCP integration with Responses API. @@ -87,57 +150,41 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """ - Returns True if any MCP tool should be handled via the litellm proxy MCP gateway. - This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/. - """ - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - return True - if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url): - return True - return False + """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in + /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" + return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) @staticmethod - def _parse_mcp_tools( + def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: + items: Final = tuple(tools or ()) + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] + other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + return gateway_tools, other_tools + + @staticmethod + async def _split_mcp_tools( tools: Iterable[Mapping[str, object]] | None, - ) -> tuple[list[ToolParam], list[Any]]: - """ - Parse tools and separate MCP tools with litellm_proxy from other tools. + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> SplitTools: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( + [ + {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} + if (name := _proxy_path_mcp_name(tool)) in served + else tool + for tool in items + ] + ) - Returns: - Tuple of (mcp_tools_with_litellm_proxy, other_tools) - """ - mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = [] - other_tools: Final[list[Any]] = [] - - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - mcp_tools_with_litellm_proxy.append(tool) - elif isinstance(server_url, str): - # Also intercept URLs like http://localhost:4000/mcp/atlassian_test - # by rewriting them to the internal litellm_proxy format. - m = _PROXY_MCP_PATH_RE.match(server_url) - if m: - rewritten = { - **tool, - "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}", - } - mcp_tools_with_litellm_proxy.append(rewritten) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - - return mcp_tools_with_litellm_proxy, other_tools + @staticmethod + async def routes_through_gateway( + tools: Iterable[Mapping[str, object]] | None, + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> tuple[bool, ...]: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) @staticmethod async def _apply_toolset_permissions( diff --git a/litellm/router.py b/litellm/router.py index 95cabfad4bd..843b916d90f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1248,7 +1248,7 @@ class Router: selector = LeastBusyLoggingHandler(router_cache=self.cache) if register_callbacks: if isinstance(litellm.input_callback, list): - litellm.input_callback.append(selector) + litellm.logging_callback_manager.add_litellm_input_callback(selector) else: litellm.input_callback = [selector] case RoutingStrategy.USAGE_BASED_ROUTING.value: @@ -4214,10 +4214,12 @@ class Router: } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - prompt_management_deployment: Final = self.get_available_deployment( + specific_deployment: Final = kwargs.pop("specific_deployment", None) + prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), + messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally + specific_deployment=specific_deployment, + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 93dddfb3d20..88ed374dd3f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -195,6 +195,30 @@ model_list: session_affinity_ttl_seconds: 300 ``` +## Custom dimensions + +Add `custom_dimensions` under `complexity_router_config` to give domain keywords or regex patterns their own weighted signal + +```yaml +custom_dimensions: + - name: internalFrameworks + weight: 0.9 + keywords: [orbitmesh, fluxgate] + - name: sqlMigration + weight: 0.7 + patterns: ['\b(create|alter|drop)\s{1,4}table\b'] +``` + +Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request + +Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one + +Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke + +Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules + +The existing heuristic-v1 tuning quota covers custom dimensions and their weights: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML or the model API; this change adds no dashboard editor + ## Usage Once configured, use the model name like any other: diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 7d4497fb6f7..c8644f52c57 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -67,6 +67,7 @@ from litellm.types.utils import ( from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section from .config import ( CALIBRATION_EXAMPLES_HEADING, + CUSTOM_PATTERN_SCAN_CHARS, DEFAULT_CLASSIFICATION_RUBRIC, DEFAULT_CODE_KEYWORDS, DEFAULT_ESCALATION_KEYWORDS, @@ -1119,6 +1120,10 @@ class ComplexityRouter(CustomLogger): self.config.custom_technical_keywords, ) self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS + self._custom_dimensions = tuple( + (dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns)) + for dimension in self.config.custom_dimensions + ) if self.config.has_custom_tiers: self.escalation_keywords: tuple[str, ...] = () elif self.config.escalation_keywords is not None: @@ -1320,6 +1325,17 @@ class ComplexityRouter(CustomLogger): score: Final = score_high if match_count >= high_threshold else score_low return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count + def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]: + if not self._custom_dimensions: + return () + scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS] + return tuple( + (DimensionScore(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight) + for dimension, patterns in self._custom_dimensions + if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords) + or any(pattern.search(scanned) is not None for pattern in patterns) + ) + def _score_multi_step(self, text: str) -> DimensionScore: """Score based on multi-step patterns.""" hits: Final = sum(1 for p in self._multi_step_patterns if p.search(text)) @@ -1415,12 +1431,13 @@ class ComplexityRouter(CustomLogger): self._score_question_complexity(prompt), ] - # Collect signals - signals: Final = [d.signal for d in dimensions if d.signal is not None] + custom_dimensions: Final = self._score_custom_dimensions(prompt, user_text) + signals: Final = [d.signal for d in (*dimensions, *(d for d, _ in custom_dimensions)) if d.signal is not None] - # Compute weighted score weights: Final = self.config.dimension_weights - weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + sum( + dimension.score * weight for dimension, weight in custom_dimensions + ) boundaries: Final = self._effective_tier_boundaries() clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score() diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index c483a0b7073..3ec9f9b5394 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -5,13 +5,21 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas All values are configurable via proxy config.yaml. """ -from collections.abc import Mapping +import math +import re +import warnings +from collections.abc import Iterable, Mapping from enum import Enum from types import MappingProxyType -from typing import Annotated, Final, Literal +from typing import Annotated, Final, Literal, NamedTuple from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator +with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + import sre_constants + import sre_parse + from litellm.types.llms.openai import REASONING_EFFORT from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin @@ -569,6 +577,117 @@ class ClassifierLLMConfig(BaseModel): return self +MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64 +MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048 +MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192 +MAX_CUSTOM_PATTERN_DEPTH: Final[int] = 16 +CUSTOM_PATTERN_SCAN_CHARS: Final[int] = 2048 + +_ATOM_OPCODES: Final = frozenset( + {sre_constants.LITERAL, sre_constants.NOT_LITERAL, sre_constants.ANY, sre_constants.IN, sre_constants.CATEGORY} +) +_REPEAT_OPCODES: Final = frozenset({sre_constants.MAX_REPEAT, sre_constants.MIN_REPEAT}) + + +class _PatternCost(NamedTuple): + paths: int + steps: int + + +def _atom_steps(node: object) -> int: + if isinstance(node, tuple) and len(node) == 2 and node[0] is sre_constants.IN: + return 1 + len(node[1]) + return 1 + + +def _repeat_cost(argument: object) -> _PatternCost | str: + if not isinstance(argument, tuple) or len(argument) != 3: + return "unsupported repeat structure" + low, high, body = argument + if high > MAX_CUSTOM_PATTERN_REPEAT or len(body) != 1 or body[0][0] not in _ATOM_OPCODES: + return "requires a single character or class repeated at most 64 times; use {n,m} instead of *, + or {n,}" + choices: Final = high - low + 1 + return _PatternCost(choices, 1 + high * _atom_steps(body[0]) + choices) + + +def _node_cost(node: object, depth: int) -> _PatternCost | str: + if not isinstance(node, tuple) or len(node) != 2: + return "unsupported regex structure" + opcode, argument = node + if opcode in _ATOM_OPCODES or opcode is sre_constants.AT: + return _PatternCost(1, _atom_steps(node)) + if opcode is sre_constants.SUBPATTERN: + return _sequence_cost(argument[-1], depth + 1) + if opcode is sre_constants.BRANCH: + costs: Final = tuple(_sequence_cost(branch, depth + 1) for branch in argument[1]) + refused: Final = next((cost for cost in costs if isinstance(cost, str)), None) + if refused is not None: + return refused + return _PatternCost( + sum(cost.paths for cost in costs if isinstance(cost, _PatternCost)), + len(costs) + sum(cost.steps for cost in costs if isinstance(cost, _PatternCost)), + ) + if opcode in _REPEAT_OPCODES: + return _repeat_cost(argument) + return "contains an unsupported regex construct" + + +def _sequence_cost(nodes: Iterable[object], depth: int) -> _PatternCost | str: + if depth > MAX_CUSTOM_PATTERN_DEPTH: + return "nests deeper than 16 levels" + costs: Final = tuple(_node_cost(node, depth) for node in nodes) + refused: Final = next((cost for cost in costs if isinstance(cost, str)), None) + if refused is not None: + return refused + valid: Final = tuple(cost for cost in costs if isinstance(cost, _PatternCost)) + # Choices multiply across a sequence; every continuation can execute once per preceding path. + total: Final = _PatternCost( + math.prod(cost.paths for cost in valid), + 1 + sum(cost.steps * math.prod(prior.paths for prior in valid[:index]) for index, cost in enumerate(valid)), + ) + if total.steps > MAX_CUSTOM_PATTERN_WORK: + return "exceeds the per-pattern regex work budget" + return total + + +def custom_pattern_work(pattern: str) -> int | str: + try: + re.compile(pattern, re.IGNORECASE) + parsed: Final = sre_parse.parse(pattern, re.IGNORECASE) + except (re.error, RecursionError, OverflowError): + return "is not a valid regex" + cost: Final = _sequence_cost(tuple(parsed), 0) + return cost if isinstance(cost, str) else cost.steps + + +class CustomDimension(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + name: str = Field(min_length=1, max_length=64, pattern=r"^[A-Za-z][A-Za-z0-9_]*$") + weight: float = Field(gt=0, le=1, allow_inf_nan=False) + keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) + patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) + + @model_validator(mode="after") + def _validate_matchers(self) -> "CustomDimension": + matchers: Final = (*self.keywords, *self.patterns) + if not matchers or any(not matcher.strip() for matcher in matchers): + raise ValueError("custom dimensions require nonblank keywords and/or patterns") + if len(matchers) > 32 or sum(map(len, matchers)) > 4096: + raise ValueError("custom dimensions allow at most 32 matchers and 4096 matcher characters each") + costs: Final = tuple((pattern, custom_pattern_work(pattern)) for pattern in self.patterns) + rejected: Final = tuple(f"pattern {pattern!r} {work}" for pattern, work in costs if isinstance(work, str)) + if rejected: + raise ValueError("custom dimension " + "; ".join(rejected)) + return self + + def pattern_work(self) -> int: + """Combined work estimate of the validated patterns.""" + return sum( + work for work in (custom_pattern_work(pattern) for pattern in self.patterns) if isinstance(work, int) + ) + + class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" @@ -671,6 +790,19 @@ class ComplexityRouterConfig(BaseModel): description="Weights for each scoring dimension", ) + custom_dimensions: tuple[CustomDimension, ...] = Field( + default=(), + max_length=16, + description=( + "Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once " + "when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. " + "Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, " + "backreferences and lookarounds are rejected. Conservative work limits include alternation paths, " + "repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. " + "Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota." + ), + ) + # Keyword lists (overridable) code_keywords: list[str] | None = Field( default=None, @@ -1245,6 +1377,27 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_custom_dimensions(self) -> "ComplexityRouterConfig": + if not self.custom_dimensions: + return self + if self.classifier_type not in ("heuristic", "heuristic_first", "hybrid"): + raise ValueError("custom_dimensions requires classifier_type heuristic, heuristic_first or hybrid") + names: Final = tuple(dimension.name.casefold() for dimension in self.custom_dimensions) + reserved: Final = frozenset(name.casefold() for name in DEFAULT_DIMENSION_WEIGHTS) + weighted: Final = frozenset(name.casefold() for name in self.dimension_weights) + if len(frozenset(names)) != len(names) or frozenset(names) & reserved: + raise ValueError("custom dimension names must be unique and must not shadow built-in dimensions") + if frozenset(names) & weighted: + raise ValueError("custom dimension weights must be inline, not in dimension_weights") + work: Final = sum(dimension.pattern_work() for dimension in self.custom_dimensions) + if work > MAX_CUSTOM_DIMENSIONS_WORK: + raise ValueError( + f"custom_dimensions regex work estimate is {work}; the limit across the router is " + f"{MAX_CUSTOM_DIMENSIONS_WORK}" + ) + return self + @field_validator("heuristic_first_max_tier", mode="before") @classmethod def _coerce_heuristic_first_max_tier(cls, value: object) -> object: diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..14e6592e1fd 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,17 +1,103 @@ -#### What this does #### -# identifies least busy deployment -# How is this achieved? -# - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"} -# - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic} -# - use litellm.success + failure callbacks to log when a request completed -# - in get_available_deployment, for a given model group name -> pick based on traffic - -import random +from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 + + +class _ModelInfo(TypedDict, total=False): + id: ReadOnly[str | int | None] + + +class _Metadata(TypedDict, total=False): + model_group: ReadOnly[str | None] + + +class _LitellmParams(TypedDict, total=False): + metadata: ReadOnly[_Metadata | None] + model_info: ReadOnly[_ModelInfo | None] + + +class _CallKwargs(TypedDict, total=False): + litellm_params: ReadOnly[_LitellmParams | None] + + +class _DeploymentModelInfo(TypedDict): + id: ReadOnly[str | int] + + +class _Deployment(TypedDict): + model_info: ReadOnly[_DeploymentModelInfo] + + +_CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) +_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) + + +def _request_count_key(model_group: str, deployment_id: str) -> str: + return f"{model_group}_request_count:{deployment_id}" + + +def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: + try: + call: Final = _CALL_KWARGS.validate_python(kwargs) + except ValidationError: + return None + litellm_params: Final = call.get("litellm_params") + metadata: Final = litellm_params.get("metadata") if litellm_params else None + model_info: Final = litellm_params.get("model_info") if litellm_params else None + model_group: Final = metadata.get("model_group") if metadata else None + deployment_id: Final = model_info.get("id") if model_info else None + if model_group is None or deployment_id is None: + return None + return model_group, str(deployment_id) + + +def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + return tuple( + _request_count_key(model_group, str(deployment["model_info"]["id"])) + for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + ) + + +def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: + return tuple(0 if value is None else int(value) for value in values) + + +def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: + values: Final = _MEMORY_COUNTS.validate_python(raw) + if values is None or len(values) != len(keys): + return (0,) * len(keys) + return _as_counts(values) + + +def _least_busy( + healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] +) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] + + +def _warn_unreadable(model_group: str, error: Exception) -> None: + verbose_router_logger.warning( + "least-busy routing could not read the shared in-flight counts for %s, " + "falling back to this worker's own counts: %s", + model_group, + error, + ) + + +def _warn_unwritable(key: str, error: Exception) -> None: + verbose_router_logger.warning("least-busy routing could not update the in-flight count under %s: %s", key, error) + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False @@ -20,195 +106,101 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache + self.router_cache_id = str(id(router_cache)) - def log_pre_api_call(self, model, messages, kwargs): - """ - Log when a model is being used. + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + self._increment(kwargs, 1) - Caching based on model group. - """ - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - request_count_api_key: Final = f"{model_group}_request_count" - # update cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + 1 + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - except Exception: - pass + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - def log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - def _get_available_deployments( - self, - healthy_deployments: list, - all_deployments: dict, - ): - """ - Helper to get deployments using least busy strategy - """ - for d in healthy_deployments: - ## if healthy deployment not yet used - if d["model_info"]["id"] not in all_deployments: - all_deployments[d["model_info"]["id"]] = 0 - # map deployment to id - # pick least busy deployment - min_traffic = float("inf") - min_deployment = None - for k, v in all_deployments.items(): - if v < min_traffic: - min_traffic = v - min_deployment = k - if min_deployment is not None: - ## check if min deployment is a string, if so, cast it to int - for m in healthy_deployments: - if m["model_info"]["id"] == min_deployment: - return m - min_deployment = random.choice(healthy_deployments) - else: - min_deployment = random.choice(healthy_deployments) - return min_deployment + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 def get_available_deployments( - self, - model_group: str, - healthy_deployments: list, - ): - """ - Sync helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys))) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) - async def async_get_available_deployments(self, model_group: str, healthy_deployments: list): - """ - Async helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + async def async_get_available_deployments( + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys))) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) + + def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + local: Final = self.router_cache.increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local < 0: + self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + if redis_cache is None: + return + redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + except Exception as e: + _warn_unwritable(key, e) + + async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + local: Final = await self.router_cache.async_increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local is not None and local < 0: + await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + if redis_cache is None: + return + await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + except Exception as e: + _warn_unwritable(key, e) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index bb6877a032b..805d4ff9080 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -32,6 +32,12 @@ def _average_latency(samples: Sequence[float]) -> float: return sum(samples) / len(samples) +def _ttft_seconds(elapsed: timedelta | float) -> float: + if isinstance(elapsed, timedelta): + return elapsed.total_seconds() + return float(elapsed) + + class LowestLatencyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -86,14 +92,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms - time_to_first_token: float | None = None total_tokens = 0 if isinstance(response_obj, ModelResponse): @@ -111,13 +116,6 @@ class LowestLatencyLoggingHandler(CustomLogger): else: final_value = response_seconds - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) - # ------------ # Update usage # ------------ @@ -138,14 +136,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -252,7 +250,7 @@ class LowestLatencyLoggingHandler(CustomLogger): {model_group}_map: { id: { "latency": [..] - "time_to_first_token": [..] + "time_to_first_token_seconds": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -273,14 +271,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms total_tokens = 0 - time_to_first_token: float | None = None if isinstance(response_obj, ModelResponse): _usage: Final = getattr(response_obj, "usage", None) @@ -296,13 +293,6 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value = float(normalized_value) else: final_value = response_seconds - - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) # ------------ # Update usage # ------------ @@ -328,14 +318,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -433,7 +423,7 @@ class LowestLatencyLoggingHandler(CustomLogger): or float("inf") ) item_latency = item_map.get("latency", []) - item_ttft_latency = item_map.get("time_to_first_token", []) + item_ttft_latency = item_map.get("time_to_first_token_seconds", []) item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py index b82269d5824..74f7b82389a 100644 --- a/litellm/router_utils/auto_router_tuning_baseline.py +++ b/litellm/router_utils/auto_router_tuning_baseline.py @@ -20,6 +20,7 @@ HEURISTIC_V1_TUNING_FIELDS: Final = ( "reasoning_override_min_score", "token_thresholds", "dimension_weights", + "custom_dimensions", "code_keywords", "reasoning_keywords", "technical_keywords", diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index 9e7f457f631..ef29f7d8fd3 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -12,6 +12,7 @@ from typing_extensions import TypedDict from litellm import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker if TYPE_CHECKING: @@ -36,10 +37,19 @@ _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0 class CooldownCache: - def __init__(self, cache: DualCache, default_cooldown_time: float): + def __init__( + self, + cache: DualCache, + default_cooldown_time: float, + redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, + ): self.cache = cache self.default_cooldown_time = default_cooldown_time self.in_memory_cache = InMemoryCache() + self._cooldown_store = DualCache( + in_memory_cache=self.in_memory_cache, + default_redis_batch_cache_expiry=redis_read_interval_seconds, + ) # Initialize the masker with custom settings for exception strings self.exception_masker = SensitiveDataMasker( visible_prefix=50, # Show first 50 characters @@ -48,6 +58,21 @@ class CooldownCache: mask_short_values=False, # Truncate long messages only; keep short ones readable ) + @property + def cooldown_store(self) -> DualCache: + """ + The cache cooldown entries live in, with the router's Redis attached on first use. + + It is kept separate from the router-wide cache so that a key missing from memory is + re-read from Redis every `redis_read_interval_seconds` rather than on the router + cache's much longer batch interval, which is what lets a sibling replica see a + cooldown another replica wrote, and so that unrelated router keys cannot evict a + cooldown from the in-memory tier before it expires. Redis is attached lazily because + the router builds its cooldown cache before it wires up the shared Redis client. + """ + self._cooldown_store.attach_redis_cache(self.cache.redis_cache) + return self._cooldown_store + def _common_add_cooldown_logic( self, model_id: str, original_exception, exception_status, cooldown_time: float ) -> tuple[str, CooldownCacheValue]: @@ -93,7 +118,7 @@ class CooldownCache: ) # Set the cache with a TTL equal to the cooldown time - self.cache.set_cache( + self.cooldown_store.set_cache( value=cooldown_data, key=cooldown_key, ttl=_cooldown_time, @@ -122,13 +147,13 @@ class CooldownCache: cooldown_cache_value: Final = CooldownCacheValue(**result) # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time if remaining <= 0: - self.cache.in_memory_cache.delete_cache(key) + self.in_memory_cache.delete_cache(key) return None - current_expiry: Final = self.cache.in_memory_cache.ttl_dict.get(key) + current_expiry: Final = self.in_memory_cache.ttl_dict.get(key) if current_expiry is not None and current_expiry > current_time + remaining + 5: corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS) - self.cache.in_memory_cache.delete_cache(key) - self.cache.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) + self.in_memory_cache.delete_cache(key) + self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) return cooldown_cache_value async def async_get_active_cooldowns( @@ -137,12 +162,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] - # Retrieve the values for the keys using mget - ## more likely to be none if no models ratelimited. So just check redis every 1s - ## each redis call adds ~100ms latency. - - ## check in memory cache first - results: Final = await self.cache.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) + results: Final = await self.cooldown_store.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = [] if results is None or all(v is None for v in results): @@ -164,7 +184,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] active_cooldowns: Final = [] current_time: Final = time.time() @@ -184,7 +204,7 @@ class CooldownCache: keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] min_cooldown_time: float | None = None # Process the results diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 0788c8db710..70362e60495 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -90,6 +90,7 @@ class PromptCachingDeploymentCheck(CustomLogger): enable_prompt_caching=( request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None ), + request_kwargs=request_kwargs, ) model_id_dict: Final = await prompt_cache.async_get_model_id( diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 2cb42ce3fac..dbaaab62d86 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,9 +1,9 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal from pydantic import BaseModel, PrivateAttr, StrictInt -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -172,6 +172,7 @@ class AugmentedAgentCard(AgentCard): class AgentObjectPermission(TypedDict, total=False): mcp_servers: list[str] | None mcp_access_groups: list[str] | None + mcp_toolsets: ReadOnly[Sequence[str] | None] mcp_tool_permissions: dict[str, list[str]] | None models: list[str] | None agents: list[str] | None diff --git a/litellm/utils.py b/litellm/utils.py index d0e11bc9551..141ec323776 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -672,11 +672,19 @@ def load_credentials_from_list(kwargs: dict): CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor") credential_name: Final = kwargs.get("litellm_credential_name") - if credential_name and litellm.credential_list: - credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name) - for key, value in credential_accessor.items(): - if key not in kwargs: - kwargs[key] = value + if not credential_name: + return + credential: Final = CredentialAccessor.find_credential(credential_name) + if credential is None: + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + credential_name, + len(litellm.credential_list), + ) + return + for key, value in credential.credential_values.items(): + if key not in kwargs: + kwargs[key] = value def get_dynamic_callbacks( diff --git a/pyproject.toml b/pyproject.toml index f4f238dd4b9..af35c77d259 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -340,6 +340,7 @@ markers = [ "asyncio: mark test as an asyncio test", "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')", "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests", + "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension", ] filterwarnings = [ # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..d63a69de76f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2956 + "limit": 2918 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1979 + "limit": 1965 }, "ANN202": { - "limit": 831 + "limit": 829 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2914 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 217 + "limit": 207 }, "S112": { "limit": 22 diff --git a/tests/llm_translation/test_bedrock_embedding.py b/tests/llm_translation/test_bedrock_embedding.py index 56baed141da..1fc05b43b23 100644 --- a/tests/llm_translation/test_bedrock_embedding.py +++ b/tests/llm_translation/test_bedrock_embedding.py @@ -1,14 +1,12 @@ import json import os -from datetime import datetime -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import Mock, patch import pytest import base64 -import httpx import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import HTTPHandler titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} @@ -394,8 +392,6 @@ def test_bedrock_embedding_uses_correct_region_when_specified(): os.environ["AWS_REGION_NAME"] = original_region_name else: os.environ.pop("AWS_REGION_NAME", None) - - def test_bedrock_embedding_region_bug_reproduction(): """ Reproduces the bug where aws_region_name is ignored when passed explicitly. @@ -458,13 +454,3 @@ def test_bedrock_embedding_region_bug_reproduction(): os.environ["AWS_REGION_NAME"] = original_region_name else: os.environ.pop("AWS_REGION_NAME", None) - - -def test_bedrock_titan_g1_text_02_model_info(): - """Test that amazon.titan-embed-g1-text-02 has correct pricing metadata""" - model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02") - assert model_info is not None, "Model info should not be None" - assert model_info["litellm_provider"] == "bedrock" - assert model_info["mode"] == "embedding" - assert model_info["input_cost_per_token"] == 1e-07 - assert model_info["max_input_tokens"] == 8192 diff --git a/tests/llm_translation/test_bedrock_embedding_pricing.py b/tests/llm_translation/test_bedrock_embedding_pricing.py deleted file mode 100644 index 099d73fed87..00000000000 --- a/tests/llm_translation/test_bedrock_embedding_pricing.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -Tests for AWS Bedrock embedding model pricing in the model cost map. - -Regression test for the Amazon Titan Text Embeddings V2 commercial price, -which was previously set 10x too high (2e-07 instead of 2e-08). -AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens -(= $0.00002 per 1K tokens = 2e-08 per token). -""" - -import importlib - - -class TestBedrockEmbeddingPricing: - """Test suite for Bedrock embedding model pricing in the cost map.""" - - def test_titan_embed_v2_commercial_input_cost(self, monkeypatch): - """Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08).""" - # Scope the local-cost-map flag to this test only, so it does not leak - # into sibling tests. monkeypatch restores the environment on teardown. - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - - import litellm.litellm_core_utils.get_model_cost_map - import litellm - - # Reload so the cost map is re-read from the local file with the flag set. - importlib.reload(litellm.litellm_core_utils.get_model_cost_map) - importlib.reload(litellm) - - model = litellm.model_cost["amazon.titan-embed-text-v2:0"] - - assert model["input_cost_per_token"] == 2e-08 - assert model["output_cost_per_token"] == 0.0 - assert model["litellm_provider"] == "bedrock" - assert model["mode"] == "embedding" diff --git a/tests/llm_translation/test_bedrock_govcloud.py b/tests/llm_translation/test_bedrock_govcloud.py index e69a95c714d..3ac1fa7cf2e 100644 --- a/tests/llm_translation/test_bedrock_govcloud.py +++ b/tests/llm_translation/test_bedrock_govcloud.py @@ -40,38 +40,6 @@ class TestBedrockGovCloudSupport: assert "us-gov-east-1" in all_regions assert "us-gov-west-1" in all_regions - def test_govcloud_models_in_model_cost(self): - """Test that GovCloud models are present in model cost configuration""" - from litellm import model_cost - - # Test Claude models in GovCloud - assert ( - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - in model_cost - ) - assert ( - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0" - in model_cost - ) - assert ( - "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost - ) - assert ( - "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost - ) - assert "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0" in model_cost - assert "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0" in model_cost - - # Test Llama models in GovCloud - assert "bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0" in model_cost - - # Test Titan models in GovCloud - assert "bedrock/us-gov-east-1/amazon.titan-text-lite-v1" in model_cost - assert "bedrock/us-gov-west-1/amazon.titan-text-lite-v1" in model_cost - def test_govcloud_model_routing(self): """Test that GovCloud models are routed correctly""" # Test Claude model routing @@ -148,135 +116,6 @@ class TestBedrockGovCloudSupport: assert not any("us-gov-east-1" in model for model in litellm.bedrock_models) assert not any("us-gov-west-1" in model for model in litellm.bedrock_models) - def test_govcloud_model_cost_properties(self): - """Test that GovCloud models have proper cost configuration""" - from litellm import model_cost - - # Check a specific GovCloud model has all required properties - govcloud_model = model_cost[ - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ] - - assert "max_tokens" in govcloud_model - assert "max_input_tokens" in govcloud_model - assert "max_output_tokens" in govcloud_model - assert "input_cost_per_token" in govcloud_model - assert "output_cost_per_token" in govcloud_model - assert govcloud_model["litellm_provider"] == "bedrock" - assert govcloud_model["mode"] == "chat" - - def test_govcloud_model_pricing_verification(self): - """Test that GovCloud models have correct pricing that differs from base models""" - from litellm import model_cost - - # Claude Haiku 4.5 commercial list pricing is under the us.* inference profile id - base_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0" - gov_east_model = ( - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ) - gov_west_model = ( - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ) - - # Verify base model pricing (us.* inference profile: $1.10/$5.50 per MTok) - base_pricing = model_cost[base_model] - assert base_pricing["input_cost_per_token"] == 1.1e-06 - assert base_pricing["output_cost_per_token"] == 5.5e-06 - - # Verify GovCloud models have different (higher) pricing - gov_east_pricing = model_cost[gov_east_model] - gov_west_pricing = model_cost[gov_west_model] - - # GovCloud models should have ~20% higher pricing than base models - assert gov_east_pricing["input_cost_per_token"] == 1.2e-06 - assert gov_east_pricing["output_cost_per_token"] == 6e-06 - assert gov_west_pricing["input_cost_per_token"] == 1.2e-06 - assert gov_west_pricing["output_cost_per_token"] == 6e-06 - - # Verify the pricing difference is approximately 20% - assert ( - abs( - gov_east_pricing["input_cost_per_token"] - / base_pricing["input_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_east_pricing["output_cost_per_token"] - / base_pricing["output_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_west_pricing["input_cost_per_token"] - / base_pricing["input_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_west_pricing["output_cost_per_token"] - / base_pricing["output_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - - # Test Claude 3 Haiku pricing - base_haiku_model = "anthropic.claude-3-haiku-20240307-v1:0" - gov_east_haiku_model = ( - "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" - ) - gov_west_haiku_model = ( - "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" - ) - - # Verify base Haiku model pricing - base_haiku_pricing = model_cost[base_haiku_model] - assert base_haiku_pricing["input_cost_per_token"] == 2.5e-07 # 0.00000025 - assert base_haiku_pricing["output_cost_per_token"] == 1.25e-06 # 0.00000125 - - # Verify GovCloud Haiku models have different (higher) pricing - gov_east_haiku_pricing = model_cost[gov_east_haiku_model] - gov_west_haiku_pricing = model_cost[gov_west_haiku_model] - - # GovCloud Haiku models should have 20% higher pricing than base models - assert ( - gov_east_haiku_pricing["input_cost_per_token"] == 3e-07 - ) # 0.0000003 (20% higher) - assert ( - gov_east_haiku_pricing["output_cost_per_token"] == 1.5e-06 - ) # 0.0000015 (20% higher) - assert ( - gov_west_haiku_pricing["input_cost_per_token"] == 3e-07 - ) # 0.0000003 (20% higher) - assert ( - gov_west_haiku_pricing["output_cost_per_token"] == 1.5e-06 - ) # 0.0000015 (20% higher) - - # Verify the pricing difference is exactly 20% - assert ( - gov_east_haiku_pricing["input_cost_per_token"] - == base_haiku_pricing["input_cost_per_token"] * 1.2 - ) - assert ( - gov_east_haiku_pricing["output_cost_per_token"] - == base_haiku_pricing["output_cost_per_token"] * 1.2 - ) - assert ( - gov_west_haiku_pricing["input_cost_per_token"] - == base_haiku_pricing["input_cost_per_token"] * 1.2 - ) - assert ( - gov_west_haiku_pricing["output_cost_per_token"] - == base_haiku_pricing["output_cost_per_token"] * 1.2 - ) - @patch("litellm.completion") def test_govcloud_completion_cost_calculation(self, mock_completion): """Test that completion requests use correct pricing for GovCloud models""" diff --git a/tests/llm_translation/test_crusoe.py b/tests/llm_translation/test_crusoe.py index 56aa4e4cd42..576428684fc 100644 --- a/tests/llm_translation/test_crusoe.py +++ b/tests/llm_translation/test_crusoe.py @@ -4,7 +4,6 @@ Tests for Crusoe provider integration import os from unittest import mock -import litellm CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1" @@ -71,38 +70,3 @@ def test_get_llm_provider_crusoe(): ) assert model == "meta-llama/Llama-3.3-70B-Instruct" assert provider == "crusoe" - - -def test_crusoe_models_configuration(): - """Test that Crusoe models are configured correctly""" - from litellm import get_model_info - - original_model_cost = litellm.model_cost - original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") - try: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - crusoe_models = [ - "crusoe/meta-llama/Llama-3.3-70B-Instruct", - "crusoe/deepseek-ai/DeepSeek-R1-0528", - "crusoe/deepseek-ai/DeepSeek-V3-0324", - "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507", - "crusoe/moonshotai/Kimi-K2-Thinking", - "crusoe/openai/gpt-oss-120b", - "crusoe/google/gemma-3-12b-it", - ] - - for model in crusoe_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - assert model_info.get("litellm_provider") == "crusoe", ( - f"{model} should have crusoe as provider" - ) - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - finally: - litellm.model_cost = original_model_cost - if original_env is None: - os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env diff --git a/tests/llm_translation/test_hyperbolic.py b/tests/llm_translation/test_hyperbolic.py index 78817fbd902..b7206e40a4e 100644 --- a/tests/llm_translation/test_hyperbolic.py +++ b/tests/llm_translation/test_hyperbolic.py @@ -1,8 +1,3 @@ -import os -from datetime import datetime -from unittest.mock import MagicMock - -import pytest import litellm @@ -69,34 +64,6 @@ def test_hyperbolic_in_provider_lists(): assert "https://api.hyperbolic.xyz/v1" in openai_compatible_endpoints -def test_hyperbolic_models_configuration(): - """Test that Hyperbolic models are properly configured""" - import json - - # Load model configuration directly from the JSON file - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path, "r") as f: - model_data = json.load(f) - - # Test a few key models - test_models = [ - "hyperbolic/deepseek-ai/DeepSeek-V3", - "hyperbolic/Qwen/Qwen2.5-Coder-32B-Instruct", - "hyperbolic/deepseek-ai/DeepSeek-R1", - ] - - for model in test_models: - assert model in model_data - model_info = model_data[model] - assert model_info["litellm_provider"] == "hyperbolic" - assert model_info["mode"] == "chat" - assert "max_tokens" in model_info - assert "input_cost_per_token" in model_info - assert "output_cost_per_token" in model_info - - def test_hyperbolic_supported_params(): """Test that supported OpenAI parameters are correctly configured""" from litellm.llms.hyperbolic.chat.transformation import HyperbolicChatConfig diff --git a/tests/llm_translation/test_lambda_ai.py b/tests/llm_translation/test_lambda_ai.py index 7ae18828d3f..edba459b352 100644 --- a/tests/llm_translation/test_lambda_ai.py +++ b/tests/llm_translation/test_lambda_ai.py @@ -8,7 +8,6 @@ from unittest import mock import pytest import litellm -from litellm import completion from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig @@ -103,48 +102,6 @@ async def test_lambda_ai_completion_call(): raise -def test_lambda_ai_models_configuration(): - """Test that Lambda AI models are configured correctly""" - from litellm import get_model_info - - # Reload model cost map to pick up local changes - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # Clear and repopulate lambda_ai_models list after reloading model_cost - litellm.lambda_ai_models = set() - litellm.add_known_models() - - # Some Lambda AI models to test - lambda_ai_models = [ - "lambda_ai/deepseek-llama3.3-70b", - "lambda_ai/hermes3-8b", - "lambda_ai/llama3.1-8b-instruct", - "lambda_ai/llama3.2-11b-vision-instruct", - "lambda_ai/qwen25-coder-32b-instruct", - ] - - for model in lambda_ai_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - assert ( - model_info.get("litellm_provider") == "lambda_ai" - ), f"{model} should have lambda_ai as provider" - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - assert ( - model_info.get("supports_function_calling") is True - ), f"{model} should support function calling" - assert ( - model_info.get("supports_system_messages") is True - ), f"{model} should support system messages" - - # Check vision support for vision models - if "vision" in model: - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - def test_lambda_ai_model_list_populated(): """Test that lambda_ai_models list is populated correctly""" # Ensure we're using local model cost map and repopulate models diff --git a/tests/llm_translation/test_morph.py b/tests/llm_translation/test_morph.py index b91d1810d38..752fb3b9083 100644 --- a/tests/llm_translation/test_morph.py +++ b/tests/llm_translation/test_morph.py @@ -68,24 +68,6 @@ def test_morph_in_provider_lists(): ) -def test_morph_model_info(): - """Test that morph models have correct configuration.""" - import litellm - - model_info = litellm.get_model_info("morph/morph-v3-large") - - assert model_info["litellm_provider"] == "morph" - assert model_info["mode"] == "chat" - assert model_info["max_tokens"] == 16000 - assert model_info["max_input_tokens"] == 16000 - assert model_info["max_output_tokens"] == 16000 - assert model_info["input_cost_per_token"] == 9e-07 # $0.9/1M tokens - assert model_info["output_cost_per_token"] == 1.9e-06 # $1.9/1M tokens - assert model_info["supports_function_calling"] is False - assert model_info["supports_vision"] is False - assert model_info["supports_system_messages"] is True - - def test_morph_supported_params(): """Test that MorphChatConfig returns correct supported parameters.""" config = MorphChatConfig() diff --git a/tests/llm_translation/test_openai_o1.py b/tests/llm_translation/test_openai_o1.py index e188a3af647..fd25e04d67d 100644 --- a/tests/llm_translation/test_openai_o1.py +++ b/tests/llm_translation/test_openai_o1.py @@ -1,15 +1,11 @@ -import json import os -from datetime import datetime -from unittest.mock import AsyncMock, patch, MagicMock +from unittest.mock import patch - -import httpx import pytest import litellm -from litellm import Choices, Message, ModelResponse +from litellm import ModelResponse from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest @@ -74,7 +70,6 @@ async def test_o1_handle_tool_calling_optional_params( - max_tokens is translated to 'max_completion_tokens' - role 'system' is translated to 'user' """ - from openai import AsyncOpenAI from litellm.utils import ProviderConfigManager from litellm.types.utils import LlmProviders @@ -186,15 +181,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest): pass -def test_o1_supports_vision(): - """Test that o1 supports vision""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - for k, v in litellm.model_cost.items(): - if k.startswith("o1") and v.get("litellm_provider") == "openai": - assert v.get("supports_vision") is True, f"{k} does not support vision" - - def test_o3_reasoning_effort(): resp = litellm.completion( model="o3-mini", diff --git a/tests/llm_translation/test_v0.py b/tests/llm_translation/test_v0.py index 95708dd855a..e96022e1e22 100644 --- a/tests/llm_translation/test_v0.py +++ b/tests/llm_translation/test_v0.py @@ -8,7 +8,6 @@ from unittest import mock import pytest import litellm -from litellm import completion from litellm.llms.v0.chat.transformation import V0ChatConfig @@ -111,33 +110,3 @@ def test_v0_supported_params(): ] assert set(supported_params) == set(expected_params) - - -def test_v0_models_configuration(): - """Test that v0 models are configured correctly""" - from litellm import get_model_info - - # Reload model cost map to pick up local changes - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # All v0 models - v0_models = ["v0/v0-1.0-md", "v0/v0-1.5-md", "v0/v0-1.5-lg"] - - for model in v0_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - # All v0 models support vision (multimodal) - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - assert ( - model_info.get("litellm_provider") == "v0" - ), f"{model} should have v0 as provider" - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - assert ( - model_info.get("supports_function_calling") is True - ), f"{model} should support function calling" - assert ( - model_info.get("supports_system_messages") is True - ), f"{model} should support system messages" diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index 2de83778f1c..38ccfd91f95 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -1,8 +1,6 @@ # What is this? ## Unit testing for the 'get_model_info()' function import os -import traceback -import json from typing import List, Dict, Any @@ -11,7 +9,7 @@ import pytest import litellm from litellm import get_model_info -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch def test_get_model_info_simple_model_name(): @@ -49,34 +47,6 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): assert model_info["input_cost_per_token"] == 0.0 -def test_get_model_info_shows_correct_supports_vision(): - info = litellm.get_model_info("gemini/gemini-2.0-flash") - print("info", info) - assert info["supports_vision"] is True - - -def test_get_model_info_shows_assistant_prefill(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - info = litellm.get_model_info("deepseek/deepseek-chat") - print("info", info) - assert info.get("supports_assistant_prefill") is True - - -def test_get_model_info_shows_supports_prompt_caching(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - info = litellm.get_model_info("deepseek/deepseek-chat") - print("info", info) - assert info.get("supports_prompt_caching") is True - - -def test_get_model_info_finetuned_models(): - info = litellm.get_model_info("ft:gpt-3.5-turbo:my-org:custom_suffix:id") - print("info", info) - assert info["input_cost_per_token"] == 0.000003 - - def test_get_model_info_gemini_pro(): info = litellm.get_model_info("gemini-2.0-flash") print("info", info) @@ -219,7 +189,7 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch): def test_get_model_info_custom_provider(): # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server: import litellm - from litellm import CustomLLM, completion, get_llm_provider + from litellm import CustomLLM, completion class MyCustomLLM(CustomLLM): def completion(self, *args, **kwargs) -> litellm.ModelResponse: diff --git a/tests/local_testing/test_least_busy_routing.py b/tests/local_testing/test_least_busy_routing.py index 18ab8bf779d..fb83d4e601f 100644 --- a/tests/local_testing/test_least_busy_routing.py +++ b/tests/local_testing/test_least_busy_routing.py @@ -33,8 +33,8 @@ def test_model_added(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"gpt-3.5-turbo_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = "gpt-3.5-turbo_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 def test_get_available_deployments(): @@ -52,8 +52,8 @@ def test_get_available_deployments(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"{model_group}_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = f"{model_group}_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 # test_get_available_deployments() @@ -104,15 +104,20 @@ async def test_router_get_available_deployments(async_test): router.leastbusy_logger.test_flag = True model_group = "azure-model" - request_count_dict = {1: 10, 2: 54, 3: 100} - cache_key = f"{model_group}_request_count" + request_count_dict = {"1": 10, "2": 54, "3": 100} + cache_keys = { + deployment_id: f"{model_group}_request_count:{deployment_id}" + for deployment_id in request_count_dict + } if async_test is True: - await router.cache.async_set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + await router.cache.async_set_cache(key=cache_keys[deployment_id], value=count) deployment = await router.async_get_available_deployment( model=model_group, messages=None, request_kwargs={} ) else: - router.cache.set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + router.cache.set_cache(key=cache_keys[deployment_id], value=count) deployment = router.get_available_deployment(model=model_group, messages=None) print(f"deployment: {deployment}") assert deployment["model_info"]["id"] == "1" @@ -124,15 +129,18 @@ async def test_router_get_available_deployments(async_test): messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - return_dict = router.cache.get_cache(key=cache_key) - # wait 2 seconds time.sleep(2) + return_dict = { + deployment_id: router.cache.get_cache(key=cache_key) + for deployment_id, cache_key in cache_keys.items() + } + assert router.leastbusy_logger.logged_success == 1 - assert return_dict[1] == 10 - assert return_dict[2] == 54 - assert return_dict[3] == 100 + assert return_dict["1"] == 10 + assert return_dict["2"] == 54 + assert return_dict["3"] == 100 ## Test with Real calls ## @@ -192,9 +200,11 @@ async def test_router_atext_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.atext_completion(model=model, prompt=prompt, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" @@ -259,8 +269,10 @@ async def test_router_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.acompletion(model=model, messages=messages, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" diff --git a/tests/local_testing/test_lowest_latency_routing.py b/tests/local_testing/test_lowest_latency_routing.py index 598b1dbcaf9..aba4500199a 100644 --- a/tests/local_testing/test_lowest_latency_routing.py +++ b/tests/local_testing/test_lowest_latency_routing.py @@ -1077,73 +1077,6 @@ async def test_latency_list_trimming_discards_oldest_entry_async(): ), f"Oldest latency {oldest_latency} should have been discarded" -def test_ttft_list_trimming_discards_oldest_entry(): - """ - The time_to_first_token list trims the oldest entry when full, matching - the behavior of the latency list. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - # TTFT is only recorded when response_obj is a ModelResponse. - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - lowest_latency_logger.log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = test_cache.get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" - - @pytest.mark.asyncio async def test_timeout_penalty_discards_oldest_entry(): """ @@ -1269,72 +1202,3 @@ def test_list_order_preserved_after_multiple_trims(): assert ( abs(latency_list[i] - expected) < tolerance ), f"At index {i}, expected ~{expected}, got {latency_list[i]}" - - -@pytest.mark.asyncio -async def test_ttft_list_trimming_discards_oldest_entry_async(): - """ - Async counterpart: the time_to_first_token list trims the oldest entry - when full. Exercises the async_log_success_event TTFT path, which only - runs when response_obj is a ModelResponse and the call is marked as - streaming with a completion_start_time. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - await lowest_latency_logger.async_log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = await test_cache.async_get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" diff --git a/tests/local_testing/test_redis_increment_with_floor.py b/tests/local_testing/test_redis_increment_with_floor.py new file mode 100644 index 00000000000..e358d5f31e0 --- /dev/null +++ b/tests/local_testing/test_redis_increment_with_floor.py @@ -0,0 +1,80 @@ +"""Least-busy routing keeps its in-flight counters in Redis, and the clamp at zero plus the +create-once TTL both live inside a Lua script. Nothing but a real Redis runs that script, so +these are the only tests that fail when the script itself is wrong.""" + +import os +import uuid +from typing import Final + +import pytest +from dotenv import load_dotenv + +load_dotenv() + +from litellm.caching.redis_cache import RedisCache + +TTL: Final = 600 + + +@pytest.fixture +def counter(): + cache: Final = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT")) + key: Final = f"lit7039-{uuid.uuid4()}" + yield cache, key, cache.check_and_fix_namespace(key=key) + cache.delete_cache(key) + + +def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter): + cache, key, _ = counter + + assert cache.increment_with_floor(key, 3, TTL) == 3 + assert cache.increment_with_floor(key, 2, TTL) == 5 + assert cache.batch_get_counts([key]) == (5,) + + +def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter): + """A worker whose counter expired mid-request decrements a key that is no longer there. + Without the clamp that deployment reads negative, and least-busy pins every later request + on it until the count climbs back to zero.""" + cache, key, _ = counter + + assert cache.increment_with_floor(key, 1, TTL) == 1 + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.batch_get_counts([key]) == (0,) + + +def test_traffic_never_pushes_a_counters_expiry_back_out(counter): + """The TTL is what releases a count whose worker died mid-request. Rewriting it on every + touch would keep that stuck count alive for as long as the group takes traffic.""" + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + assert cache.redis_client.ttl(namespaced_key) > TTL - 60 + + cache.redis_client.expire(namespaced_key, 30) + cache.increment_with_floor(key, 1, TTL) + + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter): + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + cache.redis_client.expire(namespaced_key, 30) + + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +@pytest.mark.asyncio +async def test_the_async_counter_behaves_the_same_way(counter): + cache, key, namespaced_key = counter + + assert await cache.async_increment_with_floor(key, 2, TTL) == 2 + assert await cache.async_batch_get_counts([key]) == (2,) + + cache.redis_client.expire(namespaced_key, 30) + + assert await cache.async_increment_with_floor(key, -9, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py new file mode 100644 index 00000000000..741fa7386df --- /dev/null +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -0,0 +1,264 @@ +import os +import threading +import uuid +from collections.abc import Iterator, Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager + +psycopg = pytest.importorskip("psycopg") + +pytestmark = pytest.mark.timeout(120) + +requires_db: Final = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + +HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable" +HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx" +HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC' +LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable" +LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx" +PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable" +PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx" + + +def _base_url() -> str: + return os.environ["DATABASE_URL"].split("?")[0] + + +def _index_validity(schema: str) -> Mapping[str, bool]: + with psycopg.connect(_base_url(), autocommit=True) as conn: + rows = conn.execute( + "SELECT c.relname, i.indisvalid FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_namespace n ON n.oid = c.relnamespace " + "WHERE n.nspname = %s", + (schema,), + ).fetchall() + return MappingProxyType(dict(rows)) + + +def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None: + """Abort a CONCURRENTLY build while it waits on an older snapshot, the same + spot the deadlock loser dies at, so it leaves its index INVALID.""" + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + with psycopg.connect(_base_url(), autocommit=True) as builder: + builder.execute("SET statement_timeout = '1s'") + with pytest.raises(psycopg.errors.QueryCanceled): + builder.execute(statement) + + +def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None: + _interrupt_concurrent_build( + schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})' + ) + + +def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None: + _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"') + + +@pytest.fixture +def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute( + f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)' + ) + conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)') + + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") + yield schema + + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + + +@pytest.fixture +def fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + """A brand-new database, what a first deploy sees. A scratch schema would + not do: the migrations guard on pg_constraint by name across every schema, + so a LiteLLM schema already pushed into public makes them skip and then + fail, which is exactly what CI's database looks like.""" + admin_url: Final = _base_url() + name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'CREATE DATABASE "{name}"') + + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}") + yield "public" + + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)') + + +@requires_db +def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id") + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False} + + +@requires_db +def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX) + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"') + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"') + before: Final = _index_validity(scratch_schema) + assert len(before) == 4 + assert set(before.values()) == {False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})') + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)') + conn.execute( + f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" ' + f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)' + ) + conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)') + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + +@requires_db +def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url(), autocommit=True) as other_replica: + other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)) + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"') + assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False + blocked: Final = _index_validity(scratch_schema) + assert blocked[HEALTH_INDEX] is False + assert [name for name in blocked if name.endswith("_ccnew")] + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None: + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + pinned.set() + pin.execute("SELECT pg_sleep(%s)", (seconds,)) + + +@requires_db +def test_repair_outlives_a_statement_timeout_passed_through_database_url_options( + scratch_schema: str, monkeypatch: pytest.MonkeyPatch +) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000") + pinned: Final = threading.Event() + holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0)) + holder.start() + pinned.wait() + try: + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + finally: + holder.join() + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None: + table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}" + index: Final = f"{table}_id_idx" + monkeypatch.setenv("DATABASE_URL", _base_url()) + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE public."{table}" (id TEXT)') + try: + _leave_invalid_index("public", table, index, "id") + assert _index_validity("public")[index] is False + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity("public")[index] is True + finally: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP TABLE public."{table}"') + + +def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + + +@requires_db +def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + with pytest.MonkeyPatch.context() as env: + env.setenv("DIRECT_URL", f"{_base_url()}?schema=public") + env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}") + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _invalidate_deployed_index(schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"') + _leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + +@requires_db +@pytest.mark.timeout(300) +@pytest.mark.parametrize("use_v2_resolver", [True, False]) +def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None: + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + _invalidate_deployed_index(fresh_database) + assert _index_validity(fresh_database)[HEALTH_INDEX] is False + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + + assert _index_validity(fresh_database)[HEALTH_INDEX] is True diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9a6ab08e9b6..dddc8304bfc 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1760,6 +1760,35 @@ class TestUnmanagedVertexRouting: ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") + def test_flag_on_routes_fine_tuned_endpoint_to_vertex_deployment(self): + """A fine-tuned Gemini batch stores `endpoints/` in the gs:// path; the bare model + (the endpoint id) must round-trip to the deployment configured as + `vertex_ai/gemini/` (LIT-6899).""" + endpoint_id = "7768560373388541952" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_list.return_value = [ + { + "model_name": "gemini-2.5-flash-dts-usc1", + "litellm_params": { + "model": f"vertex_ai/gemini/{endpoint_id}", + "custom_llm_provider": "vertex_ai", + }, + "model_info": {"id": "deploy-ft"}, + }, + ] + instance = self._instance(track_unmanaged=True, router=router) + job = self._job( + file_object=_unmanaged_vertex_file_object( + input_file_id=f"gs://bucket/litellm-vertex-files/endpoints/{endpoint_id}/abc.jsonl" + ) + ) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-ft", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" diff --git a/tests/router_unit_tests/test_router_cooldown_per_deployment.py b/tests/router_unit_tests/test_router_cooldown_per_deployment.py index b8ae8a8c013..228dfed38b9 100644 --- a/tests/router_unit_tests/test_router_cooldown_per_deployment.py +++ b/tests/router_unit_tests/test_router_cooldown_per_deployment.py @@ -246,12 +246,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -267,7 +267,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -290,14 +290,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -318,12 +318,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None class TestFallbackDeploymentCooldown: diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index c3edb40c819..b87f9489250 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -158,6 +158,14 @@ def test_create__vertex_ai_dispatch(seams): _assert_only(seams.vertex.create_batch, seams, "create_batch") +def test_create__vertex_ai_forwards_custom_endpoint(seams): + """The vertex handler owns the custom_endpoint batch rejection (LIT-6899), so the dispatcher + must forward the flag for the handler to act on.""" + bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True) + + assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True + + def test_create__provider_config_routes_to_base_http_handler(seams): """model + a provider batches config (bedrock-style) routes to the generic base_llm_http_handler, NOT the per-provider instance.""" diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 2f412e7382b..6b2df118611 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -525,6 +525,50 @@ def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_re assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} +def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): + """A caller that must fall back when Redis is unreachable needs the failure, not zeros. + + The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell + apart from "every counter is unset". Least-busy routing read that as an idle deployment and + kept sending traffic to it instead of falling back to this worker's own in-flight counts. + """ + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + sync_batch_redis_cache.batch_get_counts(["lit7039"]) + + +@pytest.mark.asyncio +async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None): + """Async twin: the async batch read hides the same failure behind an empty dict.""" + failing_client = AsyncMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + + with patch.object(cache, "init_async_client", return_value=failing_client): + assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + await cache.async_batch_get_counts(["lit7039"]) + + +@pytest.mark.parametrize("stored", [b"3", "3"]) +def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None): + """Counters come back positionally, so an unset key has to stay a hole rather than shift the + rest of the row onto the wrong deployments, and a count has to survive whether the client + hands it back as bytes or as text.""" + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.return_value = [stored, None, b"0"] + + assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) + + @pytest.fixture def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]: service_logger = ServiceLogging(mock_testing=True) diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index f46df5baadf..c1c3569c0e2 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks(): never sees ``data["prompt"]``.""" guardrail = _guardrail() assert guardrail.uses_apply_guardrail_interface() is True - assert guardrail._deployment_pre_call_target() is guardrail + assert guardrail._deployment_hook_target() is guardrail @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index de8b654987b..6b7780acd20 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1777,6 +1777,224 @@ class TestEnableAnthropicPromptCaching: assert messages == before +class TestClaudeCodeOneShotAutoCaching: + BILLING_TEXT = "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli; cc_is_subagent=true;" + BILLING_SYSTEM = [{"type": "text", "text": BILLING_TEXT}] + MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "unique fetched document"}]}] + + @staticmethod + def _kwargs(configured=None): + kwargs = { + "litellm_metadata": {}, + "proxy_server_request": { + "headers": { + "user-agent": "claude-cli/2.1.263 (external, cli)", + "x-app": "cli-bg", + } + }, + } + if configured is not None: + kwargs["cache_control_injection_points"] = configured + return kwargs + + @pytest.mark.parametrize( + "system", + [ + BILLING_TEXT, + BILLING_SYSTEM, + [*BILLING_SYSTEM, {"type": "text", "text": " "}], + [ + *BILLING_SYSTEM, + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli;"}, + ], + ], + ids=["string", "text_block", "whitespace_block", "multiple_billing_blocks"], + ) + @pytest.mark.parametrize("tools", [None, []], ids=["absent_tools", "empty_tools"]) + def test_skips_defaults_and_attribution_for_one_shot_subagent(self, monkeypatch, system, tools): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = copy.deepcopy(self.MESSAGES) + kwargs = self._kwargs() + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=tools, + ) + + assert result_messages == self.MESSAGES + assert result_system == system + assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + + def test_user_agent_header_lookup_is_case_insensitive(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + user_agent = kwargs["proxy_server_request"]["headers"].pop("user-agent") + kwargs["proxy_server_request"]["headers"]["User-Agent"] = user_agent + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(self.BILLING_SYSTEM), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_messages == self.MESSAGES + assert result_system == self.BILLING_SYSTEM + + def test_router_affinity_skips_string_billing_system(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = copy.deepcopy(self.MESSAGES) + kwargs = self._kwargs() + kwargs["system"] = self.BILLING_TEXT + + result = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=("claude-sonnet-4-5",), + request_kwargs=kwargs, + ) + + assert result == messages + + @pytest.mark.parametrize( + "headers,system", + [ + ("not-a-mapping", BILLING_SYSTEM), + ( + {"user-agent": "claude-cli/2.1.263 (external, cli)"}, + [{"type": "text", "text": "x-anthropic-billing-header: malformed"}], + ), + ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, None), + ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, ["not-a-mapping"]), + ( + {"user-agent": "claude-cli/2.1.263 (external, cli)"}, + [{"type": "image", "text": BILLING_TEXT}], + ), + ], + ids=["malformed_headers", "malformed_billing", "missing_system", "malformed_block", "non_text_block"], + ) + def test_malformed_untrusted_context_keeps_defaults(self, monkeypatch, headers, system): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + points = AnthropicCacheControlHook.get_default_injection_points( + messages=copy.deepcopy(self.MESSAGES), + system=copy.deepcopy(system), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + request_kwargs={"proxy_server_request": {"headers": headers}}, + ) + + assert len(points) == 2 + + def test_message_without_role_keeps_defaults(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + points = AnthropicCacheControlHook.get_default_injection_points( + messages=[{"content": "missing role"}], + system=copy.deepcopy(self.BILLING_SYSTEM), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + request_kwargs=self._kwargs(), + ) + + assert len(points) == 2 + + @pytest.mark.parametrize( + "messages,system,tools", + [ + ( + MESSAGES, + BILLING_SYSTEM, + [{"name": "WebFetch", "description": "fetch", "input_schema": {"type": "object"}}], + ), + (MESSAGES, [*BILLING_SYSTEM, {"type": "text", "text": "Explore the repository"}], None), + ( + [ + {"role": "user", "content": "first turn"}, + {"role": "assistant", "content": "reply"}, + *MESSAGES, + ], + BILLING_SYSTEM, + None, + ), + ], + ids=["tools", "real_system", "history"], + ) + def test_keeps_defaults_for_reusable_subagents(self, monkeypatch, messages, system, tools): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(messages), + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=copy.deepcopy(tools), + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2 + assert kwargs["litellm_metadata"]["litellm_gateway_injected_cache"] == "" + + @pytest.mark.parametrize( + "user_agent,system", + [ + ("anthropic-sdk-python/0.75.0", BILLING_SYSTEM), + ( + "claude-cli/2.1.263 (external, cli)", + [ + { + "type": "text", + "text": f"{BILLING_TEXT}\nadditional system instructions", + } + ], + ), + ( + "claude-cli/2.1.263 (external, cli)", + [ + { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=false;", + } + ], + ), + ], + ids=["different_client", "appended_instructions", "not_a_subagent"], + ) + def test_ambiguous_or_unmatched_signals_fail_open(self, monkeypatch, user_agent, system): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + kwargs["proxy_server_request"]["headers"]["user-agent"] = user_agent + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2 + + def test_explicit_injection_points_remain_authoritative(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs([{"location": "message", "role": "user"}]) + + result_messages, _ = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(self.BILLING_SYSTEM), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"} + + class TestPerKeyEnablePromptCaching: """Per-request enable_prompt_caching override (stamped from key metadata) with the global flag off.""" @@ -1977,6 +2195,25 @@ class TestConfiguredInjectionPointsStandDown: _, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}] + @pytest.mark.parametrize( + "configured", + [None, CONFIGURED], + ids=["automatic_defaults", "configured_points"], + ) + def test_v1_messages_stands_down_for_root_cache_control(self, monkeypatch, configured): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + root_cache_control = {"type": "ephemeral"} + kwargs = {"cache_control": root_cache_control, "litellm_metadata": {}} + if configured is not None: + kwargs["cache_control_injection_points"] = copy.deepcopy(configured) + + result_messages, result_system = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) + + assert result_messages == self.V1_MESSAGES + assert result_system == "sys" + assert kwargs["cache_control"] is root_cache_control + assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + def test_v1_messages_reentry_flow_preserves_tool_config_remainder(self): """The advisor interceptor re-enters anthropic_messages() with the outer request's kwargs and post-injection messages. The first pass applies the diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 885bd1d4d72..cd8d609cf71 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Literal, Optional +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock import pytest @@ -10,6 +10,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import CallTypes, UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail if TYPE_CHECKING: @@ -2378,11 +2379,108 @@ def _logged_call(messages: list | str) -> tuple[dict, object]: return kwargs, response +class _NativeApplyGuardrail(_InheritedApplyGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + + +@pytest.mark.parametrize("guardrail_type", (CustomGuardrail, _NativeApplyGuardrail, _InheritedApplyGuardrail)) +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), +) +def test_logging_only_requires_framework_support_or_explicit_declaration( + guardrail_type: type[CustomGuardrail], + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + supported: Final = [GuardrailEventHooks.pre_call] + if guardrail_type is _InheritedApplyGuardrail: + guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + assert guardrail.event_hook == event_hook + assert supported == [GuardrailEventHooks.pre_call] + else: + with pytest.raises(ValueError, match=r"logging_only.*not in the supported event hooks"): + guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + + explicitly_supported: Final = guardrail_type( + event_hook=event_hook, + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ) + assert explicitly_supported.event_hook == event_hook + + +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.post_call, + "post_call", + [GuardrailEventHooks.logging_only, GuardrailEventHooks.post_call], + ["logging_only", "post_call"], + Mode(tags={"enforce": "post_call"}, default="logging_only"), + Mode(tags={"enforce": ["logging_only", "post_call"]}), + Mode(tags={"audit": "logging_only"}, default="post_call"), + Mode(tags={}, default=["logging_only", "post_call"]), + ), +) +def test_framework_logging_only_does_not_allow_other_unsupported_modes( + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + with pytest.raises(ValueError, match=r"post_call.*not in the supported event hooks"): + _InheritedApplyGuardrail(event_hook=event_hook, supported_event_hooks=[GuardrailEventHooks.pre_call]) + + class TestLoggingOnlyApplyGuardrail: """LIT-4876 regression: a guardrail in mode logging_only that implements only apply_guardrail must still run against the logged request and response and record guardrail_information, instead of inheriting the CustomLogger no-op.""" + @pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), + ) + @pytest.mark.asyncio + async def test_content_filter_accepts_logging_only_and_records_detection( + self, event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode + ) -> None: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="content-review", + event_hook=event_hook, + default_on=True, + blocked_words=[BlockedWord(keyword="hello", action=ContentFilterAction.BLOCK)], + ) + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert out_response is response + assert out_kwargs["messages"] == kwargs["messages"] + assert ( + out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] + == "guardrail_intervened" + ) + @pytest.mark.asyncio async def test_runs_apply_guardrail_observe_only_and_records_verdict(self): guardrail = _ApplyOnlyObserver() @@ -2610,3 +2708,36 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: ) assert result is replacement + + @pytest.mark.asyncio + async def test_apply_guardrail_interface_modifies_deployment_response(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + class ReplacingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + assert input_type == "response" + return {**inputs, "texts": ["filtered response"]} + + guardrail = ReplacingGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.post_call, + ) + response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}]) + request_data = {"guardrails": ["test-guardrail"]} + + result = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, + response=response, + call_type=CallTypes.acompletion, + ) + + assert result is response + assert response.choices[0].message.content == "filtered response" + assert request_data == {"guardrails": ["test-guardrail"]} diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 1778eca25ef..42d3df76902 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -9,9 +9,11 @@ import litellm from litellm.litellm_core_utils.exception_mapping_utils import ( ExceptionCheckers, _get_body_error_code, + _get_response_headers, exception_type, extract_and_raise_litellm_exception, ) +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.openai.common_utils import OpenAIError from litellm.types.utils import LlmProviders @@ -1254,3 +1256,156 @@ def test_handle_error_marks_only_a_status_code_it_never_received(): raise handler._handle_error(e=upstream, provider_config=None) assert received.value.status_code == 500 assert received.value.status_code_is_synthesized is False + + +def test_bedrock_500_preserves_provider_response_headers(): + """A Bedrock 5xx must keep x-amzn-RequestId so AWS support can trace it (LIT-5428).""" + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-map-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=500, + message=provider_response.text, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-map-500" + + +@pytest.mark.parametrize( + "custom_llm_provider, status_code, provider_message, expected_exception", + [ + ( + "bedrock_mantle", + 400, + ( + '{"error":{"code":"validation_error",' + '"message":"prompt tokens (1055489) exceed model maximum (1050000) for openai.gpt-5.6-sol",' + '"param":null,"type":"invalid_request_error"}}' + ), + litellm.ContextWindowExceededError, + ), + ( + "bedrock", + 400, + '{"message":"Input is too long for requested model."}', + litellm.ContextWindowExceededError, + ), + ( + "bedrock", + 400, + '{"message":"Could not process image"}', + litellm.InternalServerError, + ), + ], +) +def test_bedrock_classified_errors_preserve_provider_response_headers( + custom_llm_provider, status_code, provider_message, expected_exception +): + """Branches that classify a Bedrock error by its text must keep x-amzn-RequestId (LIT-5428).""" + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-classified"}, + text=provider_message, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message=provider_message, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(expected_exception) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider=custom_llm_provider, + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-classified" + + +@pytest.mark.parametrize( + "status_code, provider_message", + [ + (504, '{"message":"Gateway timeout"}'), + (408, '{"message":"Bedrock did not answer in time"}'), + (408, '{"message":"Connect timeout on endpoint URL"}'), + ], +) +def test_bedrock_timeout_mapping_preserves_provider_headers(status_code, provider_message): + """A mapped bedrock timeout keeps the upstream response, like every other mapped bedrock error. + + The proxy prefixes those headers on the way out, while retry and cooldown + logic still reads the raw retry-after off the response. + """ + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-timeout", "set-cookie": "session=attacker"}, + text=provider_message, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message=provider_message, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.Timeout) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-timeout" + assert exc_info.value.headers is None + + +@pytest.mark.parametrize("status_code", [504, 408]) +def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code): + """Cooldown and retry timing read retry-after through _get_response_headers.""" + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-retry-after", "retry-after": "7"}, + text='{"message":"Bedrock did not answer in time"}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message='{"message":"Bedrock did not answer in time"}', + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.Timeout) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + exception_headers = _get_response_headers(original_exception=exc_info.value) + assert exception_headers is not None + assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7 diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py index ca25ee80c23..d24e2b58db8 100644 --- a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py +++ b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py @@ -1,6 +1,5 @@ -import litellm from litellm import LlmProviders from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -46,14 +45,6 @@ def test_xai_openai_compatible_provider_info(): assert dynamic_api_key == "api-key" -def test_xai_get_model_info_uses_xai_pricing_metadata(): - model_info = litellm.get_model_info("xai/grok-3-mini") - - assert model_info["litellm_provider"] == "xai" - assert model_info["key"] == "xai/grok-3-mini" - assert model_info["mode"] == "chat" - - def test_xai_validate_environment_reads_api_key(monkeypatch): monkeypatch.setenv("XAI_API_KEY", "api-key") diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 794613942a1..ae620fdd6dc 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -26,6 +26,30 @@ FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789" FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789" +@pytest.mark.parametrize( + "messages,system,expected", + [ + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_is_subagent=true;", True), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: =junk; cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_version=; cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: malformed", False), + ([{"content": "missing role"}], "x-anthropic-billing-header: cc_is_subagent=true;", False), + (["not-a-mapping"], "x-anthropic-billing-header: cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], ["not-a-mapping"], False), + ([{"role": "user", "content": "hi"}], None, False), + ], +) +def test_is_claude_code_one_shot_subagent_request(messages, system, expected): + from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request + + assert is_claude_code_one_shot_subagent_request( + messages=messages, + system=system, + tools=None, + user_agent="claude-cli/2.1.263 (external, cli)", + ) is expected + + class TestOptionallyHandleAnthropicOAuth: """Tests for optionally_handle_anthropic_oauth function.""" diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 2a44e77ce09..9bdc79919d2 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -1,4 +1,3 @@ -import os from unittest.mock import MagicMock import httpx @@ -38,25 +37,6 @@ class TestAzureMAIImageGeneration: assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1") - def test_mai_flash_and_2e_model_pricing_in_cost_map(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - - flash_info = litellm.get_model_info( - model="azure_ai/MAI-Image-2.5-Flash", - custom_llm_provider="azure_ai", - ) - assert flash_info["input_cost_per_token"] == 1.75e-06 - assert flash_info["input_cost_per_image_token"] == 1.75e-06 - assert flash_info["output_cost_per_image_token"] == 3.3e-05 - - image_2e_info = litellm.get_model_info( - model="azure_ai/MAI-Image-2e", - custom_llm_provider="azure_ai", - ) - assert image_2e_info["input_cost_per_token"] == 5e-06 - assert image_2e_info["output_cost_per_image_token"] == 1.95e-05 - def test_get_mai_image_generation_url(self): url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( api_base="https://my-resource.services.ai.azure.com", diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py index f3618572622..1b2ca298694 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py @@ -12,112 +12,6 @@ from importlib.resources import files import pytest -FW_MODELS = { - "azure_ai/FW-Kimi-K2.5": { - "input_cost_per_token": 6.6e-07, - "output_cost_per_token": 3.3e-06, - "cache_read_input_token_cost": 1.1e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K2.6": { - "input_cost_per_token": 1.045e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 1.76e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K2.7-Code": { - "input_cost_per_token": 1.05e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 2.1e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K3": { - "input_cost_per_token": 3.3e-06, - "output_cost_per_token": 1.65e-05, - "cache_read_input_token_cost": 3.3e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - "supports_vision": True, - }, - "azure_ai/FW-Inkling": { - "input_cost_per_token": 1e-06, - "output_cost_per_token": 4.05e-06, - "cache_read_input_token_cost": 1.7e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 1048576, - }, - "azure_ai/FW-DeepSeek-V3.2": { - "input_cost_per_token": 6.2e-07, - "output_cost_per_token": 1.85e-06, - "cache_read_input_token_cost": 3.1e-07, - "max_input_tokens": 163840, - "max_output_tokens": 163840, - }, - "azure_ai/FW-DeepSeek-V4-Pro": { - "input_cost_per_token": 1.925e-06, - "output_cost_per_token": 3.828e-06, - "cache_read_input_token_cost": 1.65e-07, - "max_input_tokens": 1000000, - "max_output_tokens": 384000, - }, - "azure_ai/FW-MiniMax-M3": { - "input_cost_per_token": 3.3e-07, - "output_cost_per_token": 1.32e-06, - "cache_read_input_token_cost": 6.6e-08, - "max_input_tokens": 512000, - "max_output_tokens": 512000, - "supports_vision": True, - }, - "azure_ai/FW-MiniMax-M2.5": { - "input_cost_per_token": 3.3e-07, - "output_cost_per_token": 1.32e-06, - "cache_read_input_token_cost": 3.3e-08, - "max_input_tokens": 1000000, - "max_output_tokens": 1000000, - }, - "azure_ai/FW-Nemotron-3-Ultra-NVFP4": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2.4e-06, - "cache_read_input_token_cost": 1.19e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - }, - "azure_ai/FW-GLM-5.2-Fast": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 6.6e-06, - "cache_read_input_token_cost": 2.1e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5.2": { - "input_cost_per_token": 1.54e-06, - "output_cost_per_token": 4.84e-06, - "cache_read_input_token_cost": 1.5e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5.1": { - "input_cost_per_token": 1.54e-06, - "output_cost_per_token": 4.84e-06, - "cache_read_input_token_cost": 2.86e-07, - "max_input_tokens": 202800, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5": { - "input_cost_per_token": 1.1e-06, - "output_cost_per_token": 3.52e-06, - "cache_read_input_token_cost": 2.2e-07, - "max_input_tokens": 200000, - "max_output_tokens": 128000, - }, -} - @pytest.fixture(scope="module") def use_local_model_cost_map(): @@ -144,28 +38,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -@pytest.mark.parametrize("model_key,expected", list(FW_MODELS.items())) -def test_azure_ai_fw_model_info(use_local_model_cost_map, model_key, expected): - model_info = use_local_model_cost_map.get_model_info(model=model_key) - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["input_cost_per_token"] == pytest.approx(expected["input_cost_per_token"]) - assert model_info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"]) - assert model_info["cache_read_input_token_cost"] == pytest.approx( - expected["cache_read_input_token_cost"] - ) - assert model_info["max_input_tokens"] == expected["max_input_tokens"] - assert model_info["max_output_tokens"] == expected["max_output_tokens"] - assert model_info["max_tokens"] == expected["max_output_tokens"] - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_prompt_caching"] is True - if expected.get("supports_vision"): - assert model_info["supports_vision"] is True - - @pytest.mark.parametrize( "model_name,expected_prompt,expected_completion", [ @@ -197,22 +69,6 @@ def test_azure_ai_fw_cost_per_token( assert completion_cost == pytest.approx(expected_completion) -def test_azure_ai_fw_nemotron_lightning_model_info(use_local_model_cost_map): - model_info = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B") - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["input_cost_per_token"] == pytest.approx(6e-08) - assert model_info["output_cost_per_token"] == pytest.approx(2.2e-07) - assert model_info["cache_read_input_token_cost"] == pytest.approx(1e-08) - assert model_info["max_input_tokens"] == 262144 - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_prompt_caching"] is True - assert model_info["supports_vision"] is False - - def test_azure_ai_fw_nemotron_lightning_supports_tool_choice(use_local_model_cost_map): from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py index 812b9288ca8..cbcc2a94043 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py @@ -33,33 +33,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -def test_azure_ai_kimi_k26_model_info(use_local_model_cost_map): - model_info = use_local_model_cost_map.get_model_info(model="azure_ai/kimi-k2.6") - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - assert model_info["max_tokens"] == 262144 - assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07) - assert model_info["output_cost_per_token"] == pytest.approx(4e-06) - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - -def test_azure_ai_kimi_k26_raw_model_cost_entry(use_local_model_cost_map): - model_info = use_local_model_cost_map.model_cost["azure_ai/kimi-k2.6"] - - assert model_info["supported_modalities"] == ["text", "image"] - assert model_info["supported_output_modalities"] == ["text"] - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - def test_azure_ai_kimi_k26_cost_per_token(use_local_model_cost_map): from litellm.llms.azure_ai.cost_calculator import cost_per_token from litellm.types.utils import Usage diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index aba51689094..c2c448cd7e2 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -177,3 +177,16 @@ def test_guardrail_config_flows_to_headers_not_request_body(model): assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "ff6ujrregl1q" assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT" assert headers["X-Amzn-Bedrock-Trace"] == "DISABLED" + + +def test_get_error_class_preserves_provider_headers(): + """The invoke handler path hands real provider headers to get_error_class (LIT-5428).""" + error = AmazonInvokeConfig().get_error_class( + error_message="Amazon Bedrock is unable to process your request.", + status_code=500, + headers={"x-amzn-RequestId": "req-invoke-500"}, + ) + + assert isinstance(error, BedrockError) + assert error.headers == {"x-amzn-RequestId": "req-invoke-500"} + assert error.response.headers["x-amzn-requestid"] == "req-invoke-500" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index cb05cdb9451..f0e361ceb88 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2,6 +2,7 @@ import asyncio import json import os +import httpx import pytest from fastapi.testclient import TestClient @@ -6039,6 +6040,8 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}} class MockResponse: + headers = httpx.Headers({"x-amzn-RequestId": "req-parse-failure"}) + def json(self): return leaky_body @@ -6067,6 +6070,7 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-parse-failure" def test_converse_drops_sampling_params_for_models_that_removed_them(): diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index 4bef59842f1..d0adabe7b4e 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -496,3 +496,171 @@ async def test_async_invoke_streaming_forwards_bedrock_response_headers(): assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987" + +def _bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_invoke_streaming_error_forwards_bedrock_response_headers(): + error_response = _bedrock_stream_error_response(500, "req-stream-err-1") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-1" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_error_forwards_bedrock_response_headers(): + error_response = _bedrock_stream_error_response(500, "req-stream-err-2") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-2" + + +def _unread_bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + stream=httpx.ByteStream(b'{"message":"Amazon Bedrock is unable to process your request."}'), + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_invoke_streaming_error_forwards_headers_when_body_was_never_read(): + """A retried streamed request raises HTTPStatusError over a body nobody read, so + reading it for the error message throws and loses the request id (LIT-5428).""" + error_response = _unread_bedrock_stream_error_response(500, "req-unread-sync") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-sync" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_error_forwards_headers_when_body_was_never_read(): + error_response = _unread_bedrock_stream_error_response(500, "req-unread-async") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-async" + + +def test_invoke_streaming_non_200_forwards_bedrock_response_headers(): + """A caller-supplied client that returns a failure instead of raising still reaches the + provider's headers, and reading the streamed body for the message must not throw (LIT-5428).""" + error_response = _unread_bedrock_stream_error_response(500, "req-non200-sync") + client = HTTPHandler() + client.post = MagicMock(return_value=error_response) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-sync" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers(): + error_response = _unread_bedrock_stream_error_response(500, "req-non200-async") + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=error_response) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-async" diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 9302dc01abe..3f03305423a 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -614,3 +614,260 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] assert "ASIABATCHSIGNROLE" in authorization assert signed_data == b'{"jobName": "litellm-batch-job"}' + + +# --------------------------------------------------------------------------- # +# Provider error headers (LIT-5428) # +# --------------------------------------------------------------------------- # + + +def _bedrock_chat_error_configs(): + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.llms.bedrock.chat.invoke_agent.transformation import AmazonInvokeAgentConfig + from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig, + ) + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, + ) + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, + ) + + return [ + AmazonInvokeConfig, + AmazonConverseConfig, + AmazonMoonshotConfig, + AmazonBedrockOpenAIConfig, + AmazonAgentCoreConfig, + AmazonInvokeAgentConfig, + ] + + +@pytest.mark.parametrize("config", _bedrock_chat_error_configs()) +def test_bedrock_chat_get_error_class_keeps_provider_headers(config): + """Every Bedrock chat route must carry x-amzn-RequestId out to the caller (LIT-5428). + + A config that drops the headers it is handed shadows the fix for its own models. + """ + error = config().get_error_class( + error_message="Amazon Bedrock is unable to process your request.", + status_code=500, + headers={"x-amzn-RequestId": "req-chat-500"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-chat-500" + + +def test_error_response_text_reads_a_read_response(): + import httpx + + from litellm.llms.bedrock.common_utils import error_response_text + + response = httpx.Response(status_code=500, text="Amazon Bedrock is unable to process your request.") + + assert error_response_text(response) == "Amazon Bedrock is unable to process your request." + + +def test_error_response_text_falls_back_when_a_streamed_response_was_never_read(): + """A retried streamed request raises HTTPStatusError over an unread body; reading it + throws ResponseNotRead and would lose the status and headers this fix preserves.""" + import httpx + + from litellm.llms.bedrock.common_utils import error_response_text + + request = httpx.Request(method="POST", url="https://bedrock-runtime.amazonaws.com") + response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-unread-500"}, + stream=httpx.ByteStream(b"never read"), + request=request, + ) + + with pytest.raises(httpx.ResponseNotRead): + _ = response.text + + assert error_response_text(response) == "Internal Server Error" + + +def test_bedrock_error_skips_header_values_httpx_cannot_carry(): + """The shared HTTP handler copies an arbitrary exception's header values in verbatim, + so a non-str value must not take down the whole error (LIT-5428).""" + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="boom", + headers={"x-amzn-RequestId": "req-mixed-500", "x-retry-count": 3, "x-nothing": None}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-mixed-500" + assert "x-retry-count" not in error.response.headers + assert isinstance(error.response, httpx.Response) + + +def test_bedrock_error_keeps_duplicate_httpx_header_values(): + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="boom", + headers=httpx.Headers([("x-amzn-RequestId", "req-dup-500"), ("set-cookie", "a=1"), ("set-cookie", "b=2")]), + ) + + assert error.response.headers.get_list("set-cookie") == ["a=1", "b=2"] + + +def _bedrock_httpx_status_error_sites(): + """Every `except httpx.HTTPStatusError as err` that raises a BedrockError, across bedrock.""" + import ast + import pathlib + + sites = [] + for path in sorted(pathlib.Path("litellm/llms/bedrock").rglob("*.py")): + tree = ast.parse(path.read_text()) + for handler in (n for n in ast.walk(tree) if isinstance(n, ast.ExceptHandler)): + caught = ast.unparse(handler.type) if handler.type is not None else "" + if "HTTPStatusError" not in caught or handler.name is None: + continue + for call in ( + n + for n in ast.walk(handler) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "BedrockError" + ): + sites.append((str(path), call.lineno, handler.name, {k.arg for k in call.keywords})) + return sites + + +def test_every_bedrock_httpx_status_error_site_keeps_provider_headers(): + """A raise site holding the provider's failed response must hand its headers on (LIT-5428). + + These sites are the only place x-amzn-RequestId still exists; a site that drops it + silently shadows the fix for that whole surface. + """ + sites = _bedrock_httpx_status_error_sites() + + assert len(sites) >= 12 + dropped = [f"{path}:{lineno}" for path, lineno, _, kwargs in sites if "headers" not in kwargs] + assert dropped == [] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_bedrock_embedding_call_keeps_provider_headers(is_async): + """The embeddings surface raises from the same shape as chat and lost the same header.""" + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + failure = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-embed-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + class _SyncUpstream(HTTPHandler): + def post(self, *args, **kwargs): + return failure + + class _AsyncUpstream(AsyncHTTPHandler): + async def post(self, *args, **kwargs): + return failure + + async def _drive(): + embedding = BedrockEmbedding() + kwargs = dict( + timeout=None, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/", + headers={}, + data={}, + ) + if is_async: + return await embedding._make_async_call(client=_AsyncUpstream(), **kwargs) + return embedding._make_sync_call(client=_SyncUpstream(), **kwargs) + + with pytest.raises(BedrockError) as exc_info: + await _drive() + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-embed-500" + + +def _bedrock_mantle_error_configs(): + from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig + from litellm.llms.bedrock_mantle.responses.transformation import BedrockMantleResponsesAPIConfig + + return [BedrockMantleChatConfig, BedrockMantleResponsesAPIConfig] + + +@pytest.mark.parametrize("config", _bedrock_mantle_error_configs()) +def test_bedrock_mantle_get_error_class_keeps_provider_headers(config): + """bedrock_mantle rides the OpenAI-compatible surfaces, whose errors drop the headers. + + A chat request for a responses-API model is bridged onto the responses config, so + fixing only the chat one leaves the model the customer actually calls uncovered. + """ + error = config().get_error_class( + error_message="prompt tokens exceed model maximum", + status_code=400, + headers={"x-amzn-RequestId": "req-mantle-400"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-mantle-400" + + +def _bedrock_configs_with_get_error_class(): + import importlib + import inspect + import pathlib + + import litellm + + llms_root = pathlib.Path(inspect.getfile(litellm)).parent / "llms" + configs = [] + for package in ("bedrock", "bedrock_mantle"): + for path in sorted((llms_root / package).rglob("*.py")): + module_name = "litellm.llms." + ".".join(path.relative_to(llms_root).with_suffix("").parts) + module = importlib.import_module(module_name) + for name, obj in vars(module).items(): + if not inspect.isclass(obj) or obj.__module__ != module_name: + continue + if getattr(obj, "get_error_class", None) is None: + continue + configs.append(pytest.param(obj, id=f"{module_name}.{name}")) + return configs + + +@pytest.mark.parametrize("config", _bedrock_configs_with_get_error_class()) +def test_every_bedrock_config_get_error_class_keeps_provider_headers(config): + """Every bedrock surface must classify errors through BedrockError, not a header-dropping base. + + A config that inherits get_error_class from a provider-agnostic base builds a blank + response, so the request id is gone before the proxy ever reads it. + """ + try: + instance = config() + except Exception: + instance = config.__new__(config) + + try: + error = instance.get_error_class( + error_message="boom", + status_code=500, + headers={"x-amzn-RequestId": "req-audit-500"}, + ) + except Exception as raised: # some bases raise the exception instead of returning it + error = raised + + assert error.response.headers["x-amzn-requestid"] == "req-audit-500" + + +def test_bedrock_get_error_class_audit_covers_every_surface(): + assert len(_bedrock_configs_with_get_error_class()) >= 30 diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 5f12ae8566c..fda3c8ceb8f 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -1,8 +1,5 @@ """Test Bedrock cross-region inference profile model mapping""" -import json -from functools import lru_cache -from pathlib import Path from typing import NamedTuple import pytest @@ -102,13 +99,6 @@ GPT_5_6_PROFILES = [ ] -@lru_cache(maxsize=1) -def _packaged_cost_map(): - """The map litellm actually resolves against, for fields ModelInfoBase drops.""" - path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json" - return json.loads(path.read_text()) - - def _bedrock_response(model, usage): return ModelResponse( id="test", @@ -126,17 +116,6 @@ def _bedrock_response(model, usage): ) -def test_bedrock_cross_region_inference_profile_mapping(): - """Test that bedrock cross-region inference profile model is mapped""" - model = "bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0" - - model_info = _get_model_info_helper(model=model, custom_llm_provider="bedrock") - - assert model_info is not None - assert model_info["litellm_provider"] == "bedrock" - assert model_info["input_cost_per_token"] == 8e-07 - - def test_proxy_cost_calculation_scenario(): """Test exact GitHub issue scenario: proxy cost calculation""" model = "litellm_proxy/bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0" @@ -176,38 +155,6 @@ def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_ma assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse" -@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) -def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map): - """Geo and Global profiles carry their own published rates, per context tier.""" - model_info = _get_model_info_helper( - model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" - ) - - assert model_info["litellm_provider"] == "bedrock_converse" - assert model_info["mode"] == "chat" - assert model_info["max_input_tokens"] == 1000000 - assert model_info["input_cost_per_token"] == profile.input_cost - assert ( - model_info["input_cost_per_token_above_272k_tokens"] - == profile.input_cost_above_272k - ) - assert model_info["output_cost_per_token"] == profile.output_cost - assert ( - model_info["output_cost_per_token_above_272k_tokens"] - == profile.output_cost_above_272k - ) - assert model_info["cache_creation_input_token_cost"] == profile.cache_write - assert ( - model_info["cache_creation_input_token_cost_above_272k_tokens"] - == profile.cache_write_above_272k - ) - assert model_info["cache_read_input_token_cost"] == profile.cache_read - assert ( - model_info["cache_read_input_token_cost_above_272k_tokens"] - == profile.cache_read_above_272k - ) - - def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map): """A prompt over 272K tokens is billed at the long-context rate, not the base rate.""" response = _bedrock_response( @@ -267,31 +214,6 @@ def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map): assert cost == pytest.approx(expected, rel=1e-9) -@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) -def test_bedrock_gpt_5_6_advertises_only_converse_supported_features( - profile, local_model_cost_map -): - model_info = _get_model_info_helper( - model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" - ) - - assert model_info["supports_function_calling"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - # Bedrock rejects an explicit cachePoint block for these models, so the flag that - # offers caller-driven caching stays off even though the cache rates are declared. - assert not model_info.get("supports_prompt_caching") - - # ModelInfoBase drops these two, so they are read from the map litellm resolves. - raw = _packaged_cost_map()[profile.model_id] - assert raw["supported_modalities"] == ["text", "image"] - assert raw["supported_output_modalities"] == ["text"] - # No bedrock_converse entry declares supported_endpoints; these models are reachable - # on chat completions and on the Responses API without it. - assert "supported_endpoints" not in raw - - @pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(profile, local_model_cost_map): """GPT-5.x on Converse maps reasoning_effort to reasoning.effort, so reasoning_effort 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 1f23d39c631..9457d5faaff 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 @@ -8,9 +8,7 @@ gate, the URL construction for both paths, and the shared Bearer auth. """ import copy -import json import logging -from pathlib import Path import pytest from botocore.exceptions import ( @@ -159,7 +157,6 @@ class TestBedrockMantleResponsesURL: assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" assert url.count("/responses") == 1 - def test_url_aws_region_name_overrides_stale_api_base(self, monkeypatch): monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) @@ -1777,53 +1774,6 @@ class TestBedrockMantleResponsesSigV4: 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") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(5.5e-06) - assert info["output_cost_per_token"] == pytest.approx(3.3e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07) - assert info["max_input_tokens"] == 1050000 - - def test_gpt_5_4_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(2.75e-06) - assert info["output_cost_per_token"] == pytest.approx(1.65e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07) - assert info["max_input_tokens"] == 1050000 - - def test_gpt_5_6_cyber_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.6-cyber") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(1.375e-05) - assert info["cache_creation_input_token_cost"] == pytest.approx(1.71875e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(1.375e-06) - assert info["output_cost_per_token"] == pytest.approx(8.25e-05) - assert info["max_input_tokens"] == 272000 - - @pytest.mark.parametrize( - "model, input_cost, cache_creation_cost, cache_read_cost, output_cost", - [ - ("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05), - ("openai.gpt-5.6-terra", 2.2e-06, 2.75e-06, 2.2e-07, 1.32e-05), - ("openai.gpt-5.6-luna", 2.2e-07, 2.75e-07, 2.2e-08, 1.32e-06), - ], - ) - def test_gpt_5_6_pricing_and_mode( - self, local_cost_map, model, input_cost, cache_creation_cost, cache_read_cost, output_cost - ): - info = litellm.get_model_info(f"bedrock_mantle/{model}") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(input_cost) - assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost) - assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost) - assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == 1050000 - assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) - assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2) - assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2) - assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5) @pytest.mark.parametrize( "model, input_cost, output_cost", @@ -1861,58 +1811,3 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models - - -def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]: - repo_root = Path(__file__).resolve().parents[4] - paths = { - "root": repo_root / "model_prices_and_context_window.json", - "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json", - } - return json.loads(paths[map_name].read_text()) - - -class TestMantleGptRegistryEntries: - """Locks the OpenAI GPT entries to Bedrock Mantle's live behavior. - - Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna - and for gpt-5.5 and gpt-5.4 (oversize requests 400 with "prompt tokens (N) - exceed model maximum (1050000)", and a 1,030,590-token request completes - on every one of them), while the AWS model cards still quote 272K for - gpt-5.5 and gpt-5.4. mode must stay "responses": Mantle's native - /v1/chat/completions rejects function tools unless reasoning_effort is - "none", so chat traffic has to keep bridging to the Responses API - (see the responses_api_bridge tests above). - """ - - @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) - @pytest.mark.parametrize( - "key", - ( - "bedrock_mantle/openai.gpt-5.6-sol", - "bedrock_mantle/openai.gpt-5.6-terra", - "bedrock_mantle/openai.gpt-5.6-luna", - ), - ) - def test_entry_matches_mantle_enforced_limits(self, map_name, key): - entry = _repo_cost_map(map_name)[key] - assert entry["max_input_tokens"] == 1050000 - assert entry["max_output_tokens"] == 128000 - assert entry["mode"] == "responses" - assert entry["use_openai_responses_path"] is True - assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] - - @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) - @pytest.mark.parametrize( - "key", - ( - "bedrock_mantle/openai.gpt-5.5", - "bedrock_mantle/openai.gpt-5.4", - ), - ) - def test_gpt_55_and_54_entries_match_mantle_enforced_limits(self, map_name, key): - entry = _repo_cost_map(map_name)[key] - assert entry["max_input_tokens"] == 1050000 - assert entry["max_output_tokens"] == 128000 - assert entry["mode"] == "responses" - assert entry["use_openai_responses_path"] is True diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index b88c27e64b9..1be94d4daa2 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -684,40 +684,6 @@ class TestBedrockMantleProviderResolution: class TestBedrockMantlePricing: """Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing.""" - def test_gpt_oss_120b_pricing(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - # Bedrock pricing: $0.15/M input, $0.60/M output - assert info["input_cost_per_token"] == pytest.approx(1.5e-7) - assert info["output_cost_per_token"] == pytest.approx(6e-7) - - def test_gpt_oss_20b_pricing(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-20b") - # Bedrock pricing: $0.075/M input, $0.30/M output - assert info["input_cost_per_token"] == pytest.approx(7.5e-8) - assert info["output_cost_per_token"] == pytest.approx(3e-7) - - def test_pricing_significantly_cheaper_than_openai_native(self, monkeypatch): - """ - Verify Bedrock Mantle pricing is cheaper than OpenAI's direct API pricing. - This is the core issue the provider addition fixes — previously users were being - billed at OpenAI rates instead of the cheaper Bedrock rates. - """ - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - bedrock_info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - # OpenAI direct pricing for gpt-oss-120b is ~$0.039/M input, $0.190/M output - # Bedrock should be cheaper at $0.15/M input and $0.60/M output... wait - # Actually, Bedrock ADDS value not reduces cost vs OpenAI direct for these models. - # The key fix is that we now use Bedrock-specific prices instead of mapping to - # some unrelated OpenAI model (like gpt-4) pricing. - # Just validate the pricing is as expected from AWS docs. - assert bedrock_info["input_cost_per_token"] == pytest.approx(1.5e-7) - assert bedrock_info["output_cost_per_token"] == pytest.approx(6e-7) - def test_safeguard_models_have_larger_output_tokens(self, monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") litellm.add_known_models() @@ -727,49 +693,6 @@ class TestBedrockMantlePricing: ) assert info_safeguard["max_output_tokens"] > info_120b["max_output_tokens"] - def test_reasoning_support(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - assert info.get("supports_reasoning") is True - - def test_context_window(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - assert info["max_input_tokens"] == 131072 - - -@pytest.mark.parametrize( - "model_id,input_cost,output_cost,max_tokens", - [ - ("google.gemma-4-31b", 1.4e-07, 4e-07, 256000), - ("google.gemma-4-26b-a4b", 1.3e-07, 4e-07, 256000), - ("google.gemma-4-e2b", 4e-08, 8e-08, 128000), - ], -) -def test_gemma_4_bedrock_mantle_model_metadata( - local_cost_map, model_id, input_cost, output_cost, max_tokens -): - full_model_name = f"bedrock_mantle/{model_id}" - info = litellm.get_model_info(full_model_name) - - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == pytest.approx(input_cost) - assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == max_tokens - assert info["max_output_tokens"] == max_tokens - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert ( - litellm.supports_parallel_function_calling( - model=full_model_name, custom_llm_provider="bedrock_mantle" - ) - is False - ) - @pytest.mark.parametrize( "model_id", diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index ca79c8d7025..12b5f03aedc 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -308,3 +308,64 @@ def test_completion_plumbs_stream_chunk_size_through_converse(): stream_chunk_size=2048, ) iter_bytes_spy.assert_called_once_with(chunk_size=2048) + + +def _bedrock_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + text=json.dumps({"message": "Amazon Bedrock is unable to process your request."}), + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_converse_completion_error_forwards_bedrock_response_headers(): + error_response = _bedrock_error_response(500, "req-err-123") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-123" + + +@pytest.mark.asyncio +async def test_async_converse_completion_error_forwards_bedrock_response_headers(): + error_response = _bedrock_error_response(500, "req-err-456") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-456" 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 e16855da8cb..98a5f4b2db5 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 @@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams @@ -37,6 +38,69 @@ from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, Trans _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "OCR output", "images": []}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} + + +def _ocr_sync_client() -> HTTPHandler: + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))) + return client + + +def _ocr_async_client() -> AsyncHTTPHandler: + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)) + ) + return client + + +def test_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = BaseLLMHTTPHandler().ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_sync_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + + +@pytest.mark.asyncio +async def test_async_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = await BaseLLMHTTPHandler().async_ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_async_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index e6fe01be4ba..d4ef4282b27 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm -from litellm import get_model_info, supports_reasoning, supports_vision +from litellm import supports_reasoning, supports_vision from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id @@ -16,17 +16,6 @@ from litellm.types.utils import ( ) -@pytest.fixture(autouse=True) -def force_local_model_cost(monkeypatch): - """Force local model cost map usage for all tests in this file.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - # Refresh model_cost from local map - import litellm - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map - - litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url) - - def test_validate_environment_sets_session_affinity_from_litellm_session_id(): config = FireworksAIConfig() @@ -404,15 +393,6 @@ def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice( assert "tool_choice" not in supported_params -def test_get_model_info_respects_explicit_fireworks_capabilities(): - """Test that get_model_info preserves explicit capability flags from the model map.""" - model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1") - - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - - def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): """Test that Fireworks only overrides supports_reasoning for supported models.""" config = FireworksAIConfig() diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py index 5641439aa54..41f6ad9d99d 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py @@ -56,17 +56,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -@pytest.mark.parametrize("alias", KIMI_ALIASES) -def test_fireworks_kimi_raw_cost_entry_limits(use_local_model_cost_map, alias): - entry = use_local_model_cost_map.model_cost[alias] - - assert entry["litellm_provider"] == "fireworks_ai" - assert entry["max_input_tokens"] == CONTEXT_WINDOW - assert entry["max_output_tokens"] == OUTPUT_LIMIT - assert entry["max_tokens"] == OUTPUT_LIMIT - assert entry["max_output_tokens"] < entry["max_input_tokens"] - - @pytest.mark.parametrize("alias", KIMI_ALIASES) def test_fireworks_kimi_get_model_info_limits(use_local_model_cost_map, alias): model_info = use_local_model_cost_map.get_model_info(model=alias) diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index 3d8200bc474..2b3b6343fad 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,13 +1,11 @@ import json -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock -import httpx import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig -from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents def test_gemini_realtime_transformation_session_created(): @@ -308,20 +306,6 @@ def test_gemini_realtime_transformation_generation_complete(): assert contains_audio_done_event, "Expected audio done event" -def test_gemini_3_1_flash_live_preview_model_cost_map_entry(): - for key in ( - "gemini-3.1-flash-live-preview", - "gemini/gemini-3.1-flash-live-preview", - ): - assert key in litellm.model_cost - info = litellm.model_cost[key] - assert "/v1/realtime" in info.get("supported_endpoints", []) - assert info.get("max_input_tokens") == 131072 - assert info.get("max_output_tokens") == 65536 - assert "video" in info.get("supported_modalities", []) - assert info.get("supports_function_calling") is True - - def test_gemini_realtime_tool_call_transformation(): """Test transformation of Gemini toolCall to OpenAI function_call_arguments.done format.""" config = GeminiRealtimeConfig() @@ -1845,19 +1829,6 @@ def test_is_audio_only_live_model_uses_cost_map(model, expected, patch_gemini_au assert GeminiRealtimeConfig._is_audio_only_live_model(model) == expected -def test_gemini_live_native_audio_entry_is_vertex_only(): - import json - from pathlib import Path - from typing import Final - - catalog_path: Final = Path(__file__).parents[5] / "model_prices_and_context_window.json" - catalog: Final = json.loads(catalog_path.read_text()) - vertex_key: Final = "gemini-live-2.5-flash-native-audio" - assert catalog[vertex_key]["litellm_provider"] == "vertex_ai-language-models" - assert catalog[vertex_key].get("gemini_native_audio") is True - assert "gemini/gemini-live-2.5-flash-native-audio" not in catalog, "the Gemini API does not serve this model" - - def test_is_setup_message_and_is_content_message(): config = GeminiRealtimeConfig() assert config.is_setup_message({"setup": {}}) is True diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py index cff3c6be940..4c0f5969249 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py @@ -231,26 +231,6 @@ def test_inception_in_provider_lists(): assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints -def test_inception_model_configuration(monkeypatch): - from litellm import get_model_info - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.inception_models = set() - litellm.add_known_models() - - info = get_model_info("inception/mercury-2") - assert info.get("litellm_provider") == "inception" - assert info.get("mode") == "chat" - assert info.get("max_input_tokens") == 128000 - assert info.get("input_cost_per_token") == 2.5e-07 - assert info.get("output_cost_per_token") == 7.5e-07 - assert info.get("cache_read_input_token_cost") == 2.5e-08 - assert info.get("supports_function_calling") is True - assert info.get("supports_tool_choice") is True - assert info.get("supports_response_schema") is True - - def test_inception_model_list_populated(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py index 62688a13c35..ed3f34fc744 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py @@ -143,24 +143,6 @@ async def test_inception_fim_async(): assert r.choices[0].text == "a + b" -def test_inception_fim_model_configuration(monkeypatch): - from litellm import get_model_info - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.text_completion_inception_models = set() - litellm.add_known_models() - - assert ( - "text-completion-inception/mercury-edit-2" - in litellm.text_completion_inception_models - ) - info = get_model_info("text-completion-inception/mercury-edit-2") - assert info.get("litellm_provider") == "text-completion-inception" - assert info.get("mode") == "completion" - assert info.get("max_input_tokens") == 32000 - - def test_inception_fim_targets_fim_endpoint(): """ End-to-end: a FIM request must hit `/v1/fim/completions` (NOT diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py index 8c8bea00dea..d484fa437ae 100644 --- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py @@ -708,38 +708,6 @@ class TestKimiK26ModelRegistry: """Load directly from the bundled backup so tests don't depend on remote fetch.""" return GetModelCostMap.load_local_model_cost_map() - def test_kimi_k26_in_model_cost_map(self, model_cost_map): - """kimi-k2.6 should be present in the model cost map.""" - assert "moonshot/kimi-k2.6" in model_cost_map, "moonshot/kimi-k2.6 not found in model_cost" - - def test_kimi_k26_pricing(self, model_cost_map): - """kimi-k2.6 pricing should match official Kimi API rates.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07) - assert model_info["output_cost_per_token"] == pytest.approx(4e-06) - assert model_info["cache_read_input_token_cost"] == pytest.approx(1.6e-07) - - def test_kimi_k26_context_window(self, model_cost_map): - """kimi-k2.6 should have a 256K (262144 token) context window.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - assert model_info["max_tokens"] == 262144 - - def test_kimi_k26_capabilities(self, model_cost_map): - """kimi-k2.6 should support function calling, vision, video input, tool choice, and reasoning.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info.get("supports_function_calling") is True - assert model_info.get("supports_tool_choice") is True - assert model_info.get("supports_vision") is True - assert model_info.get("supports_video_input") is True - assert model_info.get("supports_reasoning") is True - - def test_kimi_k26_provider(self, model_cost_map): - """kimi-k2.6 should be assigned to the moonshot provider.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["litellm_provider"] == "moonshot" - class TestMoonshotResponseSchemaSupport: """Every model currently live on api.moonshot.ai supports json_schema @@ -762,10 +730,6 @@ class TestMoonshotResponseSchemaSupport: def model_cost_map(self): return GetModelCostMap.load_local_model_cost_map() - @pytest.mark.parametrize("model", LIVE_MODELS) - def test_live_model_supports_response_schema(self, model, model_cost_map): - assert model_cost_map[model].get("supports_response_schema") is True - def test_supports_response_schema_utility_reports_true(self, model_cost_map, monkeypatch): monkeypatch.setattr(litellm, "model_cost", model_cost_map) assert litellm.utils.supports_response_schema(model="moonshot/kimi-k2.5") is True diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/test_litellm/llms/openai_like/test_cognition_provider.py index 5c71b60e08a..d392abc6cc5 100644 --- a/tests/test_litellm/llms/openai_like/test_cognition_provider.py +++ b/tests/test_litellm/llms/openai_like/test_cognition_provider.py @@ -110,22 +110,6 @@ class TestCognitionProviderIdentity: class TestCognitionCostTracking: - @pytest.mark.parametrize( - "model, input_cost, output_cost, cache_read_cost", - [ - ("cognition/swe-1.6", 5e-07, 2.5e-06, 2e-07), - ("cognition/swe-1.7", 5e-07, 2.5e-06, 2e-07), - ("cognition/swe-1.7-lightning", 2.5e-06, 1.25e-05, 1e-06), - ], - ) - def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float, cache_read_cost: float): - info = litellm.get_model_info(model=model) - - assert info["litellm_provider"] == "cognition" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cache_read_cost @pytest.mark.parametrize( "model, expected_prompt_cost, expected_completion_cost", diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py index c8743e1809d..d84cc8d3237 100644 --- a/tests/test_litellm/llms/openai_like/test_json_providers.py +++ b/tests/test_litellm/llms/openai_like/test_json_providers.py @@ -2,10 +2,9 @@ Tests for JSON-based provider configuration system. """ -import json import os import sys -from unittest.mock import MagicMock, patch +from unittest.mock import patch try: import pytest @@ -318,25 +317,6 @@ class TestDarkbloom: assert config is not None assert config.custom_llm_provider == "darkbloom" - def test_darkbloom_model_cost_map(self): - with open( - os.path.join(workspace_path, "model_prices_and_context_window.json") - ) as f: - model_cost = json.load(f) - - expected_models = { - "darkbloom/gemma-4-26b": (3e-08, 1.65e-07), - "darkbloom/gpt-oss-20b": (1.45e-08, 7e-08), - } - for model, (input_cost, output_cost) in expected_models.items(): - assert model in model_cost - assert model_cost[model]["litellm_provider"] == "darkbloom" - assert model_cost[model]["max_output_tokens"] == 32768 - assert model_cost[model]["supports_function_calling"] is True - assert model_cost[model]["supports_tool_choice"] is True - assert model_cost[model]["input_cost_per_token"] == input_cost - assert model_cost[model]["output_cost_per_token"] == output_cost - class TestPublicAIIntegration: """Integration tests for PublicAI provider""" diff --git a/tests/test_litellm/llms/openai_like/test_libertai_provider.py b/tests/test_litellm/llms/openai_like/test_libertai_provider.py index fdbe3046e9b..c17eaf7c87f 100644 --- a/tests/test_litellm/llms/openai_like/test_libertai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_libertai_provider.py @@ -59,23 +59,6 @@ class TestLibertAIProviderConfig: assert api_base == "https://custom.example.com/v1" assert api_key == "sk-test" - def test_libertai_model_cost_map(self): - """Test that libertai models are present in the model cost map""" - model_cost = litellm.model_cost - - assert "libertai/qwen3.6-27b" in model_cost - info = model_cost["libertai/qwen3.6-27b"] - assert info["litellm_provider"] == "libertai" - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - - # thinking variants are marked as reasoning models - assert ( - model_cost["libertai/qwen3.6-27b-thinking"].get("supports_reasoning") - is True - ) - def test_libertai_router_config(self): """Test that libertai can be used in Router configuration""" from litellm import Router @@ -95,20 +78,6 @@ class TestLibertAIProviderConfig: assert len(router.model_list) == 1 assert router.model_list[0]["model_name"] == "libertai-chat" - def test_libertai_model_modes(self): - """Chat models carry mode 'chat'; the embedding model carries mode 'embedding'.""" - model_cost = litellm.model_cost - - # chat model - assert model_cost["libertai/qwen3.6-27b"]["mode"] == "chat" - - # embedding model (bge-m3) must be normalized to mode 'embedding' so - # /embeddings routing and the supported-endpoints matrix stay consistent - assert "libertai/bge-m3" in model_cost - bge = model_cost["libertai/bge-m3"] - assert bge["litellm_provider"] == "libertai" - assert bge["mode"] == "embedding" - def test_libertai_supported_endpoints_matrix(self): """The runtime-served backup matrix (GET /public/supported_endpoints) lists libertai.""" import json diff --git a/tests/test_litellm/llms/openai_like/test_meta_provider.py b/tests/test_litellm/llms/openai_like/test_meta_provider.py index 11b78828da6..c79e4b77cc5 100644 --- a/tests/test_litellm/llms/openai_like/test_meta_provider.py +++ b/tests/test_litellm/llms/openai_like/test_meta_provider.py @@ -193,19 +193,6 @@ class TestMetaAnthropicMessages: class TestMuseSparkModelInfo: - def test_muse_spark_pricing_and_capabilities(self): - info = litellm.get_model_info("meta/muse-spark-1.1") - - assert info["litellm_provider"] == "meta" - assert info["input_cost_per_token"] == 1.25e-06 - assert info["output_cost_per_token"] == 4.25e-06 - assert info["cache_read_input_token_cost"] == 1.5e-07 - assert info["max_input_tokens"] == 1048576 - assert info["supports_reasoning"] is True - assert info["supports_web_search"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True def test_muse_spark_cost_calculation(self): from litellm import completion_cost diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py index 6a6271e95e2..6ca7072e7ab 100644 --- a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py +++ b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py @@ -3,7 +3,6 @@ Unit tests for Perplexity embedding transformation logic. """ import base64 -import json import struct from unittest.mock import MagicMock @@ -298,25 +297,3 @@ class TestPerplexityEmbeddingProviderConfig: ) assert config is not None assert isinstance(config, PerplexityEmbeddingConfig) - - -class TestPerplexityEmbeddingModelInfo: - """Test that Perplexity embedding models are in model_prices_and_context_window.""" - - def test_model_info_available(self): - import litellm - - info = litellm.get_model_info("perplexity/pplx-embed-v1-0.6b") - assert info is not None - assert info["mode"] == "embedding" - assert info["max_input_tokens"] == 32768 - assert info["output_vector_size"] == 1024 - - def test_model_info_4b_available(self): - import litellm - - info = litellm.get_model_info("perplexity/pplx-embed-v1-4b") - assert info is not None - assert info["mode"] == "embedding" - assert info["max_input_tokens"] == 32768 - assert info["output_vector_size"] == 2560 diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py index 6630039e92e..7556b215e66 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py @@ -26,7 +26,6 @@ from litellm.types.utils import ( Usage, PromptTokensDetailsWrapper, ) -from litellm.utils import get_model_info class TestPerplexityCostCalculator: @@ -317,21 +316,6 @@ class TestPerplexityCostCalculator: assert math.isclose(total_cost, expected_total, rel_tol=1e-6) - def test_model_info_access(self): - """Test that model info correctly returns the new cost fields.""" - model_info = get_model_info( - model="sonar-deep-research", custom_llm_provider="perplexity" - ) - - # Check that the new fields are accessible - assert "citation_cost_per_token" in model_info - assert model_info["citation_cost_per_token"] == 2e-6 - assert model_info["search_context_cost_per_query"] == { - "search_context_size_low": 0.005, - "search_context_size_medium": 0.005, - "search_context_size_high": 0.005, - } - @pytest.mark.parametrize("citation_tokens", [0, 10, 25, 100]) @pytest.mark.parametrize("search_queries", [0, 1, 5, 10]) @pytest.mark.parametrize("reasoning_tokens", [0, 15, 30]) @@ -477,37 +461,6 @@ class TestPerplexityCostCalculator: assert math.isclose(prompt_cost, expected_prompt, rel_tol=1e-9) assert math.isclose(completion_cost, expected_completion, rel_tol=1e-9) - @pytest.mark.parametrize( - "model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read", - [ - ("deepseek-v4-flash-0731", 0.13, 0.26, 0.028), - ("glm-5.2", 1.4, 4.4, 0.14), - ("kimi-k3", 3.0, 15.0, 0.3), - ("kimi-k2.7-code", 0.95, 4.0, 0.19), - ], - ) - def test_agent_api_entries_carry_perplexity_published_rates( - self, model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read - ): - """The Agent API third-party models are priced from Perplexity's own catalog - (GET https://api.perplexity.ai/v1/models, `pricing` in usd_per_1m_tokens). - Perplexity's model id already starts with `perplexity/`, so the cost-map key - doubles the prefix. Regression: glm-5.2 shipped glm-5.3's 0.26 cache-read rate, - copied from the neighbouring catalog row, an 86% overcharge on cached input. - """ - info = get_model_info( - model=f"perplexity/{model_id}", custom_llm_provider="perplexity" - ) - - assert info["key"] == f"perplexity/perplexity/{model_id}" - assert info["litellm_provider"] == "perplexity" - assert info["mode"] == "responses" - assert math.isclose(info["input_cost_per_token"], usd_per_1m_input / 1e6, rel_tol=1e-9) - assert math.isclose(info["output_cost_per_token"], usd_per_1m_output / 1e6, rel_tol=1e-9) - assert math.isclose( - info["cache_read_input_token_cost"], usd_per_1m_cache_read / 1e6, rel_tol=1e-9 - ) - def test_agent_api_fallback_rates_price_a_response_without_metered_cost(self): """Perplexity meters cost on the response, but when `usage.cost` is absent the calculator falls back to the mapped per-token rates. Regression: that fallback diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 38fde3caa63..24df3214da3 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -35,7 +35,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest - from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) @@ -178,6 +177,197 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): sync_client.post.assert_not_called() +def test_create_batch_sync_does_not_resolve_publisher_models(): + """Publisher-model jobs must not incur the endpoint-resolution GET, and the job model must + stay the publisher path untouched.""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001" + safe_get.assert_not_called() + + +ENDPOINT_ID = "7768560373388541952" +ENDPOINT_CREATE_DATA = { + "input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid" +} +TUNED_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/1234509876" + + +def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = { + "name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + "deployedModels": ( + deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}] + ), + } + return resp + + +def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): + """A fine-tuned Gemini file id must produce a batch job against the endpoint's deployed + tuned model resource; the v1 batch API rejects endpoint resources in `model` (LIT-6899).""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()) as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + get_args, get_kwargs = safe_get.call_args + assert get_args[1] == ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}" + ) + assert get_kwargs["headers"]["Authorization"] == f"Bearer {TOKEN}" + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == TUNED_MODEL_RESOURCE + + +@pytest.mark.parametrize( + "api_base, expected", + [ + ( + None, + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/v1", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/vertex", + f"https://proxy.internal/vertex/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ], +) +def test_build_endpoint_resolution_url(api_base, expected): + """A custom api_base must replace the Google host for the endpoint-resolution GET without + producing a malformed url (no ':' grafting, no doubled /v1).""" + url = VertexAIBatchPrediction._build_endpoint_resolution_url( + api_base=api_base, + model=f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + vertex_location=LOCATION, + ) + assert url == expected + + +def test_create_batch_sync_endpoint_resolution_error_raises(): + h = _make_handler() + client = MagicMock() + resolve_response = MagicMock() + resolve_response.status_code = 404 + resolve_response.text = "endpoint not found" + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=resolve_response), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 404 + client.post.assert_not_called() + + +def test_create_batch_custom_endpoint_raises_400_without_io(): + """custom_endpoint deployments have no Vertex batch surface; creating a job would target a + nonexistent publisher model, so the handler must 400 before any auth or HTTP work (LIT-6899).""" + h = _make_handler() + client = MagicMock() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert exc_info.value.status_code == 400 + assert "custom_endpoint" in str(exc_info.value) + h._ensure_access_token.assert_not_called() + client.post.assert_not_called() + + +def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): + h = _make_handler() + client = MagicMock() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response(deployed_models=[])), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "no deployed model" in str(exc_info.value) + client.post.assert_not_called() + + def test_create_batch_sync_httpstatuserror_propagates(): """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the sync create path must surface that error, not swallow it.""" diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index ccb2d7e310d..232c6413e78 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -34,6 +34,12 @@ INPUT_FILE = ( "models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) +ENDPOINT_ID = "7768560373388541952" +ENDPOINT_INPUT_FILE = ( + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" + "e9412502-2c91-42a6-8e61-f5c294cc0fc8" +) + # =========================================================================== # # transform_openai_batch_request_to_vertex_ai_batch_request @@ -67,6 +73,41 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +def test_transform_openai_request_fine_tuned_endpoint_builds_endpoint_resource(): + """A fine-tuned Gemini file id (endpoints/) must target the endpoint resource, + not a nonexistent publisher model (LIT-6899).""" + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + vertex_location="us-central1", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_defaults_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_without_project_raises_400(): + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": ENDPOINT_INPUT_FILE}) + assert exc_info.value.status_code == 400 + assert "vertex_project" in str(exc_info.value) + + +def test_transform_openai_request_publisher_model_ignores_project_and_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": INPUT_FILE}, + vertex_project="my-project", + vertex_location="europe-west4", + ) + assert job["model"] == "publishers/google/models/gemini-1.5-flash-001" + + @pytest.mark.parametrize( "input_file_id", [ @@ -321,6 +362,35 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): assert exc_info.value.status_code == 400 +def test_get_model_from_gcs_file_fine_tuned_endpoint(): + """The whole endpoint id must survive parsing; the old 3-segment publishers/ parse dropped it.""" + assert T._get_model_from_gcs_file(ENDPOINT_INPUT_FILE) == f"endpoints/{ENDPOINT_ID}" + + +def test_get_model_from_gcs_file_publisher_path_wins_over_endpoints_prefix(): + """A bucket prefix containing endpoints/ must not override the publisher model path + LiteLLM appended after it.""" + uri = "gs://bucket/team-endpoints/999/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/uuid" + assert T._get_model_from_gcs_file(uri) == "publishers/google/models/gemini-1.5-flash-001" + + +def test_get_model_from_gcs_file_last_endpoints_segment_wins(): + """With no publisher path, the endpoint id closest to the file (last occurrence) is the one + LiteLLM stored; an earlier prefix segment must not shadow it.""" + uri = f"gs://bucket/endpoints/999/litellm-vertex-files/endpoints/{ENDPOINT_ID}/uuid" + assert T._get_model_from_gcs_file(uri) == f"endpoints/{ENDPOINT_ID}" + + +def test_get_model_from_gcs_file_non_numeric_endpoints_segment_raises_400(): + with pytest.raises(VertexAIError) as exc_info: + T._get_model_from_gcs_file("gs://bucket/endpoints/not-a-number/file-uuid") + assert exc_info.value.status_code == 400 + + +def test_get_bare_model_name_from_gcs_file_fine_tuned_endpoint(): + assert T.get_bare_model_name_from_gcs_file(ENDPOINT_INPUT_FILE) == ENDPOINT_ID + + # =========================================================================== # # is_unmanaged_gcs_batch_input_file_id # =========================================================================== # @@ -334,6 +404,8 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): ("file-abc123", False), ("gs://bucket/no-model-here.jsonl", False), ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + (ENDPOINT_INPUT_FILE, True), + ("gs://bucket/endpoints/not-a-number/file-uuid", False), ], ) def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 3c2d56997b7..8a249820cbd 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -159,6 +159,91 @@ class TestCreateFileUrl: assert "?" not in object_name +class TestBatchObjectNaming: + def test_should_store_publisher_model_under_publishers_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}]) + assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/") + + def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config): + """A numeric endpoint id must not be filed under publishers/google/models/gemini/, + which the batch transformation later mangles into a nonexistent publisher model (LIT-6899).""" + object_name = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {"model": "gemini/7768560373388541952"}}] + ) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "publishers" not in object_name + + def test_should_store_bare_numeric_endpoint_under_endpoints_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "7768560373388541952"}}]) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + + def test_deployment_model_overrides_jsonl_body_model(self, config): + """The stored path decides which Vertex model the batch later runs against with the + deployment's credentials, so a user-crafted JSONL body.model must not be able to redirect + an authorized deployment to a different endpoint.""" + object_name = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {"model": "9999999999999999999"}}], + deployment_model="vertex_ai/gemini/7768560373388541952", + ) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "9999999999999999999" not in object_name + + def test_url_derives_object_path_from_configured_model(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={ + "gcs_bucket_name": "my-bucket", + "model": "vertex_ai/gemini/7768560373388541952", + }, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "9999999999999999999"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "9999999999999999999" not in object_name + + +class TestCustomEndpointBatchUpload: + def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config): + """custom_endpoint deployments have no Vertex batch surface; the upload must 400 instead + of staging a file that can only produce a doomed batch job (LIT-6899).""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + with pytest.raises(VertexAIError) as exc_info: + config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "openai/gemma-2-2b-it"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + assert exc_info.value.status_code == 400 + assert "custom_endpoint" in str(exc_info.value) + + def test_should_allow_non_batch_upload_for_custom_endpoint_deployment(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("notes.txt", b"hello", "text/plain"), + "purpose": "assistants", + }, + ) + assert "/b/my-bucket/" in url + + class TestTransformRetrieveFile: def test_should_build_correct_gcs_metadata_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 6ba8706b0d8..04e46eab1b7 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -136,19 +136,6 @@ class TestVertexAIVideoConfig: # Should NOT include endpoint assert not url.endswith(":predictLongRunning") - def test_veo_31_lite_model_cost_entries_match_pricing(self): - for path in (ROOT_MODEL_COST_PATH, BACKUP_MODEL_COST_PATH): - model_cost = _load_model_cost_map(path) - info = model_cost.get(VEO_31_LITE_VERTEX_MODEL) - - assert info is not None, f"{VEO_31_LITE_VERTEX_MODEL} missing from {path}" - assert info["litellm_provider"] == "vertex_ai-video-models" - assert info["mode"] == "video_generation" - assert info["max_input_tokens"] == 1024 - assert info["output_cost_per_second"] == 0.05 - assert info["output_cost_per_second_1080p"] == 0.08 - assert info["supported_modalities"] == ["text", "image"] - def test_veo_31_lite_provider_routing_from_local_model_map( self, monkeypatch: pytest.MonkeyPatch ): diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py index 83c3bf1ecef..e591c1ae682 100644 --- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py +++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py @@ -63,11 +63,6 @@ TIER_COST_FIELDS = ( "output_cost_per_token_above_200k_tokens", "cache_read_input_token_cost_above_200k_tokens", ) -STALE_TIER_FIELDS = ( - "input_cost_per_token_above_128k_tokens", - "output_cost_per_token_above_128k_tokens", - "cache_read_input_token_cost_above_128k_tokens", -) def expected_retirement_date(slug: str) -> str: @@ -102,11 +97,10 @@ def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str): assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug) -@pytest.mark.parametrize("slug", REDIRECTED_SLUGS) -def test_no_slug_keeps_the_superseded_128k_tier(cost_map: dict, slug: str): - """The 128k tier belonged to the retired model; grok-4.3 tiers at 200k.""" - for field in STALE_TIER_FIELDS: - assert field not in cost_map[slug], field +def test_a_live_xai_model_is_untouched(cost_map: dict): + """Guard against the repricing leaking onto models xAI still serves directly.""" + assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] + assert "deprecation_date" not in cost_map["xai/grok-4.6"] @pytest.mark.parametrize("slug", REDIRECTED_SLUGS) @@ -116,12 +110,7 @@ def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str entry = cost_map[slug] for field in TIER_COST_FIELDS: assert entry[field] == target[field], field - - -def test_a_live_xai_model_is_untouched(cost_map: dict): - """Guard against the repricing leaking onto models xAI still serves directly.""" - assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] - assert "deprecation_date" not in cost_map["xai/grok-4.6"] + assert {k for k in entry if "_above_" in k} == {k for k in target if "_above_" in k} def test_both_cost_maps_agree_on_the_redirected_slugs(): diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index 38ddac8d510..069ac5727f6 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -2,11 +2,9 @@ Tests for Z.AI (Zhipu AI) provider - GLM models """ -import json import math import pytest -import respx import litellm from litellm import completion @@ -57,32 +55,9 @@ def test_zai_in_provider_lists(): assert "zai" in litellm.provider_list -def test_zai_models_in_model_cost(local_model_cost_map): - """Test that ZAI models are in the model cost map""" - - zai_models = [ - "zai/glm-4.7", - "zai/glm-4.6", - "zai/glm-4.5", - "zai/glm-4.5v", - "zai/glm-4.5-x", - "zai/glm-4.5-air", - "zai/glm-4.5-airx", - "zai/glm-4-32b-0414-128k", - "zai/glm-4.5-flash", - ] - - for model in zai_models: - assert model in litellm.model_cost, f"Model {model} not found in model_cost" - assert litellm.model_cost[model]["litellm_provider"] == "zai" - - def test_zai_glm46_cost_calculation(local_model_cost_map): """Test the cost calculation for glm-4.6""" - key = "zai/glm-4.6" - info = litellm.model_cost[key] - prompt_cost, completion_cost = cost_per_token( model="zai/glm-4.6", prompt_tokens=1000000, # 1M tokens @@ -94,26 +69,6 @@ def test_zai_glm46_cost_calculation(local_model_cost_map): assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) -def test_zai_flash_model_is_free(local_model_cost_map): - """Test that glm-4.5-flash has zero cost""" - - key = "zai/glm-4.5-flash" - info = litellm.model_cost[key] - - assert info["input_cost_per_token"] == 0 - assert info["output_cost_per_token"] == 0 - - -def test_glm47_supports_reasoning(local_model_cost_map): - """Test that GLM-4.7 supports reasoning""" - - key = "zai/glm-4.7" - assert key in litellm.model_cost, f"Model {key} not found in model_cost" - - info = litellm.model_cost[key] - assert info["supports_reasoning"] is True - - def test_glm47_cost_calculation(local_model_cost_map): """Test cost calculation for GLM-4.7""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 44eb7795659..20f4719e4bf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -11,6 +11,7 @@ from starlette.datastructures import Headers from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + UnloadableEntitlementError, _is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( @@ -4396,6 +4397,147 @@ class TestAgentMCPPermissions: ) assert sorted(result) == ["tool_a", "tool_b"] + def _agent_object_permission(self, *, toolset_ids, servers=(), tool_permissions=None): + agent_object_permission = MagicMock() + agent_object_permission.mcp_servers = list(servers) + agent_object_permission.mcp_access_groups = [] + agent_object_permission.mcp_tool_permissions = tool_permissions + agent_object_permission.mcp_toolsets = list(toolset_ids) + return agent_object_permission + + def _mock_manager_with_toolsets(self, toolset_perms): + mock_manager = MagicMock() + mock_manager.expand_permission_list = MagicMock(side_effect=lambda servers: list(servers)) + mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {}) + mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value=toolset_perms) + return mock_manager + + def _agent_toolset_patches(self, agent_object_permission, mock_manager): + return ( + patch.object( # test-quality-ok: stub the agent perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, "_get_agent_object_permission", AsyncMock(return_value=agent_object_permission) + ), + patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling toolset tests + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here + MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[]) + ), + ) + + async def test_get_allowed_mcp_servers_for_agent_includes_toolset_servers(self): + """An agent granted only mcp_toolsets reaches the toolset's servers, exactly as a + key, team, or org granted only toolsets does""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"], servers=["server-direct"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + result = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + + assert sorted(result) == ["server-a", "server-direct"] + mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"]) + + async def test_get_allowed_mcp_servers_toolset_only_agent_caps_key_servers(self): + """Regression: an agent whose only grant is a toolset used to resolve to [] and place + no ceiling at all, so a key bound to it kept every server the key itself granted""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + stack.enter_context( + patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[]) + ) + ) + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == ["server-a"] + + async def test_get_allowed_mcp_servers_agent_dangling_toolset_denies(self): + """An agent toolset that resolves to nothing is a known restriction with unknown + contents: deny, never fall through to the key's own servers""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-gone"]) + mock_manager = self._mock_manager_with_toolsets({}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + with pytest.raises(UnloadableEntitlementError): + await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + stack.enter_context( + patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[]) + ) + ) + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == [] + + async def test_get_agent_tool_permissions_for_server_unions_direct_and_toolset_tools(self): + """The agent's tool ceiling on a server is its direct tool grants plus the tools its + toolsets grant there, and None only when neither names the server""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission( + toolset_ids=["toolset-1"], tool_permissions={"server-a": ["tool_direct"]} + ) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_via_toolset"], "server-b": ["tool_b"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth) + server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth) + server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth) + + assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"] + assert server_b_tools == ["tool_b"] + assert server_c_tools is None + + async def test_get_allowed_tools_for_server_toolset_only_agent_caps_key_tools(self): + """Regression: a key allowing [tool_a, tool_b] bound to an agent whose toolset grants + only tool_a on the server ends with [tool_a]; the toolset used to be ignored""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_a"]}) + key_perm = MagicMock() + key_perm.mcp_tool_permissions = {"server-a": ["tool_a", "tool_b"]} + key_perm.mcp_toolsets = [] + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + stack.enter_context( + patch.object( # test-quality-ok: stub the key perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, "_get_key_object_permission", return_value=key_perm + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it absent here + MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None) + ) + ) + result = await MCPRequestHandler.get_allowed_tools_for_server("server-a", user_api_key_auth) + + assert result == ["tool_a"] + async def test_get_agent_object_permission_uses_shared_helper(self): """``_get_agent_object_permission`` must resolve the agent's ``object_permission_id`` and then defer to the shared diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index be4206a1faf..763200c3709 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4,6 +4,7 @@ import hashlib import json import time from base64 import urlsafe_b64encode +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -18,6 +19,57 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None): + credential = {"type": "oauth2", "access_token": access_token} + if refresh_token is not None: + credential["refresh_token"] = refresh_token + if expires_in_seconds is not None: + credential["expires_at"] = (datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)).isoformat() + if expires_at is not None: + credential["expires_at"] = expires_at + return credential + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fields", "egress_has_token"), + [ + (None, False), + ({"access_token": "", "refresh_token": "refresh-token"}, False), + ({}, True), + ({"expires_at": "never"}, True), + ({"expires_in_seconds": 600}, True), + ({"expires_in_seconds": 30}, False), + ({"expires_in_seconds": -300}, False), + ({"expires_in_seconds": -300, "refresh_token": ""}, False), + ({"expires_in_seconds": 30, "refresh_token": "refresh-token"}, True), + ({"expires_in_seconds": -300, "refresh_token": "refresh-token"}, True), + ], +) +async def test_vendor_credential_state_agrees_with_egress_token_resolution(monkeypatch, fields, egress_has_token): + from litellm.proxy._experimental.mcp_server import db as mcp_db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + + monkeypatch.setattr(mcp_db, "MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", 60) + credential = _stored_grant(**fields) if fields is not None else None + read = AsyncMock(return_value=credential) + refresh = AsyncMock(return_value=_stored_grant(access_token="fresh-token", expires_in_seconds=3600)) + prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr(mcp_db, "get_user_oauth_credential", read) + monkeypatch.setattr(mcp_db, "refresh_user_oauth_token", refresh) + + connect = await discoverable_endpoints._vendor_credential_state("user-1", "server-1") + read.assert_awaited_once_with(prisma, "user-1", "server-1") + refresh.assert_not_awaited() + egress = await mcp_db.resolve_valid_user_oauth_token( + user_id="user-1", server=MagicMock(), cred=credential, prisma_client=prisma + ) + + assert (egress is not None) is egress_has_token + assert connect == ("present" if egress_has_token else "absent") + + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @pytest.fixture(autouse=True) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1670370f082..73a52a8d2e8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -230,7 +231,7 @@ async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_co assert location.path == "/ui/connect" params = parse_qs(location.query) handle = params["connect_flow"][0] - assert params["connect_client"] == ["https://claude.ai"] + assert set(params) == {"connect_flow"} set_cookie = response.headers["set-cookie"] assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie assert "HttpOnly" in set_cookie @@ -889,14 +890,60 @@ def _opened_principal(payload): return admitted.principal -async def _finish_connect_page(response): +class _VendorCredential: + def __init__(self, state="present"): + self.calls = [] + self.state = state + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.state + + +class _ServerReachability: + def __init__(self, reachable=True): + self.calls = [] + self.reachable = reachable + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.reachable + + +async def _complete_page(response, scoped_server=None, vendor=None, reachable=None, cache=None, **overrides): + from unittest.mock import patch + handle, cookies = _flow_cookie_from(response) - completed = await complete_connect_flow( - request=_request("/authorize/complete", cookies=cookies, method="POST"), - flow_handle=handle, - session_user_id="u1", - cache=DualCache(), - ) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache or DualCache(), + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + **overrides, + ) + + +async def _describe_page(response, scoped_server=None, vendor=None, reachable=None, session_user_id="u1", cookies=None): + from unittest.mock import patch + + handle, flow_cookies = _flow_cookie_from(response) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await describe_connect_flow( + request=_request("/authorize/flow", cookies=flow_cookies if cookies is None else cookies), + flow_handle=handle, + session_user_id=session_user_id, + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + ) + + +async def _finish_connect_page(response, scoped_server=None): + completed = await _complete_page(response, scoped_server=scoped_server) return parse_qs(urlparse(completed.headers["location"]).query)["code"][0] @@ -910,23 +957,43 @@ def _sealed_wire_json(sealed, prefix, debug_key): @pytest.mark.asyncio async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): - """LIT-4917: a per-server RFC 8707 resource naming a gateway-managed oauth2 server - seals that server into the flow. The connect page interlude runs exactly as before - (the scope restricts, it never skips consent), and the code minted at the finish step - and the session pair it redeems for are both scoped.""" + """LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2 + server seals that server into the flow. The connect URL carries only the handle; the page + learns the scoped server and its vendor state from describe_connect_flow, and the finish + step refuses to mint a scoped code until that vendor credential exists, without burning + the flow. The code minted afterwards and the session pair it redeems for are both scoped.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert ( _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id" ) - code = await _finish_connect_page(response) + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent")) + assert json.loads(described.body) == { + "state": "interactive", + "client_origin": "https://claude.ai", + "server_id": "github-id", + "server_name": "github", + "connected": False, + } + cache = DualCache() + premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache) + assert premature.status_code == 400 + assert "authorize the requested MCP server" in json.loads(premature.body)["error_description"] + present = _VendorCredential("present") + completed = await _complete_page(response, scoped_server=github, vendor=present, cache=cache) + assert completed.status_code == 303 + assert present.calls == [("u1", "github-id")] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert ( _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"] == "github-id" @@ -952,9 +1019,11 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): ) async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, resolves): """Every resource shape outside 'exactly one gateway-managed server' keeps today's flow: - connect page interlude, and NONE of the minted artifacts carry the scope key on the - wire, not the flow cookie, not the code, not the session JWT, so an unscoped flow - started on a new pod completes on a pod whose strict models predate the claim.""" + the generic connect grid (describe names no server, the finish step never consults the + vendor credential), and NONE of the minted + artifacts carry the scope key on the wire, not the flow cookie, not the code, not the + session JWT, so an unscoped flow started on a new pod completes on a pod whose strict + models predate the claim.""" import base64 from unittest.mock import patch @@ -963,10 +1032,18 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server() response = _scoped_authorize(client_id, resource) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") - code = await _finish_connect_page(response) + vendor = _VendorCredential("absent") + described = await _describe_page(response, vendor=vendor) + assert json.loads(described.body)["state"] == "unscoped" + assert json.loads(described.body)["server_id"] is None + completed = await _complete_page(response, vendor=vendor) + assert vendor.calls == [] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code") token_response = await _redeem(code, client_id) payload = json.loads(token_response.body) @@ -979,19 +1056,150 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, @pytest.mark.asyncio async def test_scoped_authorize_delegate_server_stays_unscoped(): """A delegate-auth oauth2 server is outside the gateway-managed set (its keyless flow is - upstream PKCE via the relay), so a resource naming it never scopes the gateway flow.""" + upstream PKCE via the relay), so a resource naming it never scopes the gateway flow and + never narrows the connect page to it.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True) response = _scoped_authorize(client_id, SCOPED_RESOURCE) - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} + assert json.loads((await _describe_page(response)).body)["server_id"] is None code = await _finish_connect_page(response) token_response = await _redeem(code, client_id) assert _opened_principal(json.loads(token_response.body)).resource_server_id is None +@pytest.mark.asyncio +async def test_m2m_scoped_flow_mints_without_a_user_credential(): + """A client-credentials server is already authorized by its gateway service credential, so + a resource-scoped flow finishes without consulting the per-user vault.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + m2m = _scoped_mcp_server(oauth2_flow="client_credentials") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = m2m + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + vendor = _VendorCredential("unavailable") + described = await _describe_page(response, scoped_server=m2m, vendor=vendor) + assert json.loads(described.body)["state"] == "m2m" + assert json.loads(described.body)["connected"] is True + assert vendor.calls == [] + completed = await _complete_page(response, scoped_server=m2m, vendor=vendor) + assert completed.status_code == 303 + assert vendor.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("oauth2_flow", ["authorization_code", "client_credentials"]) +async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow): + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server(oauth2_flow=oauth2_flow) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + reachable = _ServerReachability(False) + vendor = _VendorCredential("present") + described = await _describe_page(response, scoped_server=server, reachable=reachable, vendor=vendor) + assert json.loads(described.body) == { + "state": "stale", + "client_origin": "https://claude.ai", + "server_id": None, + "server_name": None, + "connected": None, + } + cache = DualCache() + refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache) + assert refused.status_code == 400 + assert vendor.calls == [] + assert reachable.calls == [("u1", "github-id"), ("u1", "github-id")] + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + + +@pytest.mark.asyncio +async def test_stale_scoped_flow_remains_distinct_from_unscoped(): + """A server removed after authorize stays a stale scoped flow, so the page cannot offer a + broader unscoped grant or report a misleading Finish action.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + described = await _describe_page(response, scoped_server=None) + assert json.loads(described.body)["state"] == "stale" + assert json.loads(described.body)["connected"] is None + stale = await _complete_page(response, scoped_server=None) + assert stale.status_code == 400 + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + + +@pytest.mark.asyncio +async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credential(): + """Cancel is the escape hatch: a scoped user who cannot finish the vendor step still ends + the flow with access_denied and no credential lookup. A scoped server that is no longer + gateway-managed refuses to mint (nothing could serve that code) but also burns nothing.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = github + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + cache = DualCache() + skipped_reachability = _ServerReachability(False) + stale = await _complete_page(response, scoped_server=None, reachable=skipped_reachability, cache=cache) + assert stale.status_code == 400 + assert skipped_reachability.calls == [] + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("unavailable")) + assert described.status_code == 503 + vendor = _VendorCredential("absent") + deny_reachability = _ServerReachability(False) + denied = await _complete_page( + response, + scoped_server=github, + vendor=vendor, + reachable=deny_reachability, + cache=cache, + decision="deny", + ) + assert denied.status_code == 303 + assert parse_qs(urlparse(denied.headers["location"]).query)["error"] == ["access_denied"] + assert vendor.calls == [] + assert deny_reachability.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "session_user_id, cookies, expected_status, expected_error", + [ + ("u1", {}, 400, "invalid_request"), + ("u1", {"mcp_connect_flow_wrong": "garbage"}, 400, "invalid_request"), + (None, None, 401, "login_required"), + ("u2", None, 403, "access_denied"), + ], +) +async def test_describe_connect_flow_refuses_exactly_like_the_finish_step( + session_user_id, cookies, expected_status, expected_error +): + """The page's read of the flow is gated the same way minting is: the HttpOnly cookie for + that handle must open and the signed-in user must be the sealed one. A lure link with a + made-up handle therefore learns nothing and starts nothing.""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize(client_id, session_user_id="u1") + described = await _describe_page(response, session_user_id=session_user_id, cookies=cookies) + assert described.status_code == expected_status + assert json.loads(described.body)["error"] == expected_error + + @pytest.mark.asyncio async def test_token_rejects_resource_conflicting_with_sealed_scope(): """RFC 8707 section 2.2: redeeming a scoped code (or rotating a scoped refresh token) @@ -1005,7 +1213,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope(): with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) - code = await _finish_connect_page(response) + code = await _finish_connect_page(response, scoped_server=github) with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = linear diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 086ab854e36..79a2a27bb6d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2,6 +2,7 @@ import asyncio import contextvars import os from datetime import datetime, timedelta +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1368,6 +1369,295 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): mock_logger.info.assert_any_call("Successfully fetched %s tools total from all MCP servers", 0) +def _denied_scope_manager(known_server_names_to_ids: dict[str, str]) -> MagicMock: + servers = {name: MagicMock(server_id=server_id) for name, server_id in known_server_names_to_ids.items()} + manager = MagicMock() + manager.get_mcp_server_by_name = lambda name, client_ip=None: servers.get(name) + return manager + + +def _scope_resolver(resolved_without_agent: dict[str, str], access_groups: tuple[str, ...] = ()) -> AsyncMock: + async def resolve(user_api_key_auth, mcp_servers, client_ip=None): + if user_api_key_auth is not None and user_api_key_auth.agent_id: + return [] + return [ + SimpleNamespace( + server_id=server_id, + server_name=server_name, + alias=None, + short_prefix=None, + access_groups=list(access_groups), + ) + for server_name, server_id in resolved_without_agent.items() + ] + + return AsyncMock(side_effect=resolve) + + +async def _denied_scoped_list( + user_api_key_auth: UserAPIKeyAuth, + mcp_servers: list[str], + mock_manager: MagicMock, + resolver: AsyncMock, +) -> HTTPException: + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + + with ( + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + resolver, + ), + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + mock_manager, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=None, + mcp_servers=mcp_servers, + ) + return exc_info.value + + +@pytest.mark.asyncio +async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent(): + """The agent-binding veto must raise a 403 naming the agent, never a silent 200 with no tools.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + resolver = _scope_resolver(resolved_without_agent={"github": "srv-github"}) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "MCP server 'github'" in message + assert "agent 'agent-123'" in message + assert "mcp_servers" in message + rerun_auth = resolver.await_args_list[1].kwargs["user_api_key_auth"] + assert rerun_auth.agent_id is None + assert rerun_auth.user_id == "test_user" + assert resolver.await_args_list[1].kwargs["mcp_servers"] == ["github"] + + +@pytest.mark.asyncio +async def test_empty_scope_lists_nothing_instead_of_raising_a_nameless_denial(): + """An empty ``x-mcp-servers`` header scopes to no servers; that is an empty listing, not a 403.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + resolver = AsyncMock(return_value=[]) + + with ( + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + resolver, + ), + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + _denied_scope_manager({"github": "srv-github"}), + ), + ): + listing = await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=[] + ) + + assert listing.tools == [] + resolver.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scoped_list_denied_for_non_agent_key_raises_generic_403(): + """A denial for a key with no agent binding stays generic and skips the agent-stripped rerun.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + resolver = AsyncMock(return_value=[]) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "github" in message + assert "agent" not in message + resolver.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scoped_list_unknown_name_raises_same_generic_403_as_unauthorized(): + """Unknown and registered-but-unauthorized names raise byte-identical generic 403s, so a + caller cannot probe which server names exist.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + unknown = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({}), _scope_resolver(resolved_without_agent={}) + ) + unauthorized = await _denied_scoped_list( + user_api_key_auth, + ["github"], + _denied_scope_manager({"github": "srv-github"}), + _scope_resolver(resolved_without_agent={}), + ) + + assert unknown.status_code == unauthorized.status_code == 403 + assert unknown.detail["error"] == unauthorized.detail["error"] + assert "github" in unknown.detail["error"] + assert "agent" not in unknown.detail["error"] + + +@pytest.mark.asyncio +async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group(): + """An access-group scope vetoed by the agent binding raises the 403 naming the agent and the + group instead of the silent empty list.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["prod-group"], + _denied_scope_manager({}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "access group 'prod-group'" in message + assert "agent 'agent-123'" in message + assert "mcp_access_groups" in message + + +@pytest.mark.asyncio +async def test_scoped_list_mixed_unknown_and_vetoed_group_names_the_group_that_resolved(): + """With an unknown name ahead of the agent-vetoed group in the scope, the 403 must name the group + whose servers the key can reach, never the unknown name, or the admin is told to grant a group + that does not exist.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["no-such-group", "prod-group"], + _denied_scope_manager({}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "access group 'prod-group'" in message + assert "no-such-group" not in message + assert "agent 'agent-123'" in message + + +@pytest.mark.asyncio +async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403(): + """When the agent-stripped rerun still resolves nothing, the 403 stays generic instead of + blaming the agent binding.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + resolver = _scope_resolver(resolved_without_agent={}) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "github" in message + assert "agent" not in message + assert resolver.await_count == 2 + + +@pytest.mark.asyncio +async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_name(): + """The scope filter matches `/mcp/GitHub` to a server named `github` case-insensitively, so the + agent-attributed 403 must match the same way instead of falling back to the generic denial.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["GitHub"], + _denied_scope_manager({"github": "srv-github"}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "MCP server 'GitHub'" in message + assert "agent 'agent-123'" in message + + +@pytest.mark.asyncio +async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(): + """The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error + (McpError, INVALID_REQUEST) carrying the denial message, instead of a raw 500.""" + try: + from litellm.proxy._experimental.mcp_server.server import handle_list_tools + except ImportError: + pytest.skip("MCP server not available") + + from mcp.shared.exceptions import McpError + from mcp.types import INVALID_REQUEST + + denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'" + denial = HTTPException(status_code=403, detail={"error": denial_message}) + + with ( + patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( # test-quality-ok: the listing helper is the handler's only collaborator; the suite's seam + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + new=AsyncMock(side_effect=denial), + ), + ): + with pytest.raises(McpError) as exc_info: + await handle_list_tools() + + assert exc_info.value.error.code == INVALID_REQUEST + assert exc_info.value.error.message == denial_message + + +@pytest.mark.asyncio +async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(): + try: + from litellm.proxy._experimental.mcp_server.server import mcp_server_tool_call + except ImportError: + pytest.skip("MCP server not available") + + denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'" + denial = HTTPException(status_code=403, detail={"error": denial_message}) + + with ( + patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( # test-quality-ok: the tool-call helper is the handler's only collaborator; the suite's seam + "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + new=AsyncMock(side_effect=denial), + ), + ): + result = await mcp_server_tool_call("github-search_issues", {}) + + assert result.isError is True + assert result.content[0].text == f"Error: {denial_message}" + + @pytest.mark.asyncio async def test_mcp_server_tool_call_body_with_none_arguments(): """Test that proxy_server_request body handles None arguments correctly""" @@ -3518,6 +3808,35 @@ async def test_call_mcp_tool_user_unauthorized_access(): assert "User not allowed to call this tool" in exc_info.value.detail +@pytest.mark.asyncio +async def test_call_mcp_tool_scoped_denial_names_the_binding_agent(): + from litellm.proxy._experimental.mcp_server.server import call_mcp_tool + + agent_bound_key = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-123") + + with ( + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_allowed_mcp_servers", + AsyncMock(return_value=[]), + ), + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + _scope_resolver({"github": "srv-github"}), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await call_mcp_tool( + name="github-search_issues", + arguments={}, + user_api_key_auth=agent_bound_key, + mcp_servers=["github"], + ) + + assert exc_info.value.status_code == 403 + assert "MCP server 'github'" in exc_info.value.detail["error"] + assert "agent 'agent-123'" in exc_info.value.detail["error"] + + @pytest.mark.asyncio async def test_call_mcp_tool_unauthorized_403_does_not_leak_server_credentials(): """Regression for LIT-4703 / GH #29936. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 798b0001af1..b8935d07774 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -11,6 +11,7 @@ Covers: import json from collections.abc import Sequence +from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -1271,3 +1272,37 @@ class TestMcpServerToolCallErrorHandling: assert result.isError is True assert "User not allowed to call this tool" in result.content[0].text + + +@pytest.mark.asyncio +async def test_handle_mcp_tool_call_scoped_denial_names_the_binding_agent() -> None: + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.tool_search import handle_mcp_tool_call + + agent_bound_key = UserAPIKeyAuth(api_key="test_key", agent_id="agent-123") + + async def resolve(user_api_key_auth, mcp_servers, client_ip=None): + if user_api_key_auth.agent_id: + return [] + return [ + SimpleNamespace( + server_id="srv-github", server_name="github", alias=None, short_prefix=None, access_groups=[] + ) + ] + + with patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(side_effect=resolve), + ): + with pytest.raises(HTTPException) as exc_info: + await handle_mcp_tool_call( + tool_name="github-create_issue", + arguments={}, + user_api_key_dict=agent_bound_key, + mcp_servers=["github"], + ) + + assert exc_info.value.status_code == 403 + assert "MCP server 'github'" in exc_info.value.detail["error"] + assert "agent 'agent-123'" in exc_info.value.detail["error"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index bf0df17fafb..f0b4e94f72f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2216,3 +2216,26 @@ async def test_top_k_above_router_default_is_respected(): assert len(filtered) == 6 print("✅ Configured top_k above the semantic-router default of 5 is honored") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves(): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + external_tool = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "allowed_tools": ["zapier_send_email"], + } + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + narrowed = await SemanticToolFilterHook._narrow_mcp_references( + [gateway_reference, external_tool], ["srv-tool_1"], served_names=served_names + ) + + assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index a996de4d40c..aaf630ad29b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3,6 +3,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext """ import base64 +import logging from typing import Optional from unittest.mock import MagicMock, patch @@ -15,6 +16,7 @@ from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, check_complete_credentials, custom_auth_common_checks_warning, + log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, get_end_user_id_from_request_body, get_key_mcp_rpm_limit, @@ -101,6 +103,41 @@ class TestWarnOnceIfCustomAuthSkipsCommonChecks: assert logger.warning.call_count == 0 +class TestLogOnceIfBudgetReservationDisabled: + @pytest.fixture(autouse=True) + def _reset_sentinel(self, monkeypatch): + monkeypatch.setattr( + "litellm.constants.budget_reservation_disabled_info_emitted", + False, + ) + + def test_logs_info_only_once_when_enabled(self, caplog): + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + log_once_if_budget_reservation_disabled(disabled=False) + assert not any( + "disable_budget_reservation is enabled" in record.message + for record in caplog.records + ) + for _ in range(3): + log_once_if_budget_reservation_disabled(disabled=True) + + records = [ + record + for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert len(records) == 1 + assert records[0].levelno == logging.INFO + + def test_logs_to_injected_logger_only_once(self): + logger = MagicMock() + log_once_if_budget_reservation_disabled(disabled=False, logger=logger) + for _ in range(3): + log_once_if_budget_reservation_disabled(disabled=True, logger=logger) + assert logger.info.call_count == 1 + assert "disable_budget_reservation is enabled" in logger.info.call_args[0][0] + + class TestGetKeyModelRpmLimit: """Tests for get_key_model_rpm_limit function.""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index d44f96d95bf..541aeabcbcd 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1,5 +1,6 @@ import asyncio import json +import logging from contextlib import contextmanager from datetime import datetime, timedelta from types import SimpleNamespace @@ -146,6 +147,35 @@ async def test_disable_budget_reservation_skips_reservation(): assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_disable_budget_reservation_does_not_log_per_request(caplog): + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + for _ in range(3): + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={"disable_budget_reservation": True}, + ) + + records = [ + record + for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert records == [] + assert user_api_key_auth_obj.budget_reservation is None + + @pytest.mark.asyncio async def test_budget_reservation_runs_when_not_disabled(): """Control for #27639: with the flag absent, the reservation still runs and is stored.""" diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index 3bb3f75b9b8..2a0ebf9492a 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -824,7 +824,7 @@ class _StopFailingSubscriber(ConfigSyncSubscriber): raise RuntimeError("stop failed") -async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None: +async def test_proxy_config_subscriber_resyncs_deployments_only() -> None: from litellm.proxy.proxy_server import ProxyConfig cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) @@ -852,10 +852,7 @@ async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> await callback() await config.stop_config_sync_subscriber() - assert calls == [ - ("add_deployment", prisma_client, proxy_logging_obj), - ("get_credentials", prisma_client, None), - ] + assert calls == [("add_deployment", prisma_client, proxy_logging_obj)] assert config.config_sync_subscriber is None assert subscriber._task is None diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index f0fbdea4e85..6f7c20166c5 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -393,6 +393,51 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc assert not proxy_server.MODEL_RECONCILE_LOCK.locked() +@pytest.mark.asyncio +async def test_resync_model_deployments_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import AsyncMock, MagicMock + + import litellm + import litellm.proxy.proxy_server as proxy_server + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + from litellm.types.utils import CredentialItem + + rows: Final = [MagicMock()] + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows) + router: Final = MagicMock() + router.get_model_list.return_value = [] + installed: Final = MagicMock() + + async def load_credentials_from_db(prisma_client: object) -> None: + CredentialAccessor.upsert_credentials( + [ + CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db"}, + credential_info={}, + ) + ] + ) + + def install_models(db_models: object) -> None: + installed(db_models=db_models, credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", load_credentials_from_db) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install_models) + + assert await _resync_model_deployments("model-created-on-a-sibling-replica") is True + installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"}) + + @pytest.mark.asyncio async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch): from unittest.mock import AsyncMock, MagicMock diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py index ca4d62737b6..54e0aa74a25 100644 --- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py @@ -23,7 +23,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import ( ) OPERATOR_CRON_JOB_ID = "spend_log_cleanup_job" -SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "get_credentials_job", "add_deployment_job") +SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "add_deployment_job") async def _noop() -> None: ... diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py index 6f67b91ac1d..03d7ea81257 100644 --- a/tests/test_litellm/proxy/db/conftest.py +++ b/tests/test_litellm/proxy/db/conftest.py @@ -1,5 +1,11 @@ +import json import os +import signal +import sys +import time from collections.abc import Generator +from dataclasses import dataclass +from pathlib import Path from typing import Optional import pytest @@ -75,3 +81,76 @@ def reset_entra_token_provider_cache() -> Generator[None, None, None]: def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DATABASE_URL", "about-to-be-unset") monkeypatch.delenv("DATABASE_URL") + + +FAKE_PRISMA_CLI = """#!{python} +import json +import os +import pathlib +import subprocess +import sys +import time + +calls_file = pathlib.Path(os.environ["FAKE_PRISMA_CALLS"]) +earlier_calls = calls_file.read_text().splitlines() if calls_file.exists() else [] +with calls_file.open("a") as log: + print(json.dumps(sys.argv[1:]), file=log) +if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"): + grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"]) + pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid)) + time.sleep(600) +sys.exit(0) +""" + + +@dataclass(frozen=True, slots=True) +class FakePrismaCli: + """A stand-in `prisma` on PATH, recording every invocation. + + With FAKE_PRISMA_HANG_FIRST set it hangs on its first call from a process tree + of its own, the way the real CLI wraps Node around a Rust schema engine, so a + timeout that kills only the direct child leaves the rest of that tree running. + """ + + calls_file: Path + grandchild_pidfile: Path + + @property + def calls(self) -> list[list[str]]: + if not self.calls_file.exists(): + return [] + return [json.loads(line) for line in self.calls_file.read_text().splitlines()] + + def grandchild_is_gone(self, within_seconds: float) -> bool: + deadline = time.monotonic() + within_seconds + while time.monotonic() < deadline: + try: + os.kill(int(self.grandchild_pidfile.read_text()), 0) + except ProcessLookupError: + return True + time.sleep(0.05) + return False + + +@pytest.fixture +def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[FakePrismaCli, None, None]: + bin_dir = tmp_path / "fakebin" + bin_dir.mkdir() + script = bin_dir / "prisma" + script.write_text(FAKE_PRISMA_CLI.format(python=sys.executable)) + script.chmod(0o755) + cli = FakePrismaCli( + calls_file=tmp_path / "calls.jsonl", + grandchild_pidfile=tmp_path / "grandchild.pid", + ) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") + monkeypatch.setenv("FAKE_PRISMA_CALLS", str(cli.calls_file)) + monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(cli.grandchild_pidfile)) + monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1") + monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False) + yield cli + if cli.grandchild_pidfile.exists(): + try: + os.kill(int(cli.grandchild_pidfile.read_text()), signal.SIGKILL) + except ProcessLookupError: + pass diff --git a/tests/test_litellm/proxy/db/test_check_migration.py b/tests/test_litellm/proxy/db/test_check_migration.py index 9e2f6a1089c..74a6841cb23 100644 --- a/tests/test_litellm/proxy/db/test_check_migration.py +++ b/tests/test_litellm/proxy/db/test_check_migration.py @@ -37,3 +37,35 @@ def test_check_migration_out_of_sync(mocker): check_migration.verbose_logger.exception.assert_called_once() actual_message = check_migration.verbose_logger.exception.call_args[0][0] assert "prisma schema out of sync with db" in actual_message + + +@pytest.mark.timeout(30) +def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch): + """ + `prisma migrate diff` ran unbounded, so a database that never answers hung boot + before uvicorn ever started, and interrupting the proxy orphaned the schema engine. + """ + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + assert fake_prisma_cli.calls == [ + ["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x", + "--to-schema-datamodel", "./schema.prisma", "--script"] + ] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch): + """ + Boot calls this helper directly, so an ImportError here takes the proxy down before + uvicorn starts. An install without the runner must lose the diagnostic, not the proxy. + """ + import sys + + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index f0983d6bf62..7f43e483557 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -10,7 +10,7 @@ from fastapi.testclient import TestClient -from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema +from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema @pytest.fixture(autouse=True) @@ -193,7 +193,10 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client( mock_new_prisma.connect.assert_awaited_once() -def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): +DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"] + + +def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url): """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call into the opt-in REPLICA IDENTITY FULL step.""" from litellm.proxy.db.prisma_client import PrismaManager @@ -208,14 +211,13 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): staticmethod(lambda: applied.append(True)), ) - with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] assert applied == [True] -def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): +def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): """A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the primary key back to ("request_id"), which Postgres rejects; the guard must fail fast with guidance instead of running the push.""" @@ -228,29 +230,23 @@ def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - with pytest.raises(RuntimeError) as err: - PrismaManager.setup_database(use_migrate=False) + with pytest.raises(RuntimeError) as err: + PrismaManager.setup_database(use_migrate=False) assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR - mock_run.assert_not_called() + assert fake_prisma_cli.calls == [] -def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch): +def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): from litellm.proxy.db.prisma_client import PrismaManager from litellm_proxy_extras.utils import ProxyExtrasDBManager monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] def _entra_jwt(expires_in_seconds: int) -> str: @@ -377,3 +373,30 @@ def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch with pytest.raises(RuntimeError, match="DATABASE_HOST"): wrapper.get_rds_iam_token() + + +@pytest.mark.timeout(45) +def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch): + """ + A timed-out `db push` used to leave Node and the schema engine writing the schema, + so the next attempt pushed into a database the abandoned one was still mutating. + """ + monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False) + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert PrismaManager.setup_database(use_migrate=False) is True + assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot( + fake_prisma_cli, unset_database_url, monkeypatch +): + """ + An ImportError out of setup_database escapes the caller's RuntimeError handler and + kills boot, bypassing the operator's enforce_prisma_migration_check choice. + """ + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert PrismaManager.setup_database(use_migrate=False) is False + assert fake_prisma_cli.calls == [] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 0dbd4591ac9..87a1b84acc5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -3,6 +3,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ import json +import logging import re from unittest.mock import patch @@ -520,6 +521,80 @@ class TestToolPermissionGuardrail: ) assert excinfo.value.status_code == 400 + @pytest.mark.asyncio + async def test_async_pre_call_hook_without_tools_logs_skip_at_debug(self, caplog): + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + result = await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + assert result is data + skip_levels = [r.levelno for r in caplog.records if "No tools or functions in data" in r.getMessage()] + assert skip_levels == [logging.DEBUG] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_pre_call_hook_denied_tool_logs_at_info(self, caplog): + data = {"tools": [{"type": "function", "function": {"name": "Read"}}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException): + await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_post_call_success_hook_denied_tool_logs_at_info(self, caplog): + tool_call = {"function": {"name": "Read", "arguments": "{}"}, "type": "function"} + response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})]) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self.guardrail.async_post_call_success_hook( + data={"guardrails": ["test-tool-permission"]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + def test_parse_tool_call_arguments_malformed_json_logs_warning(self, caplog): + tool_call = ChatCompletionMessageToolCall(function={"name": "Bash", "arguments": "{not json"}, id="call_1") + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + parsed, error = self.guardrail._parse_tool_call_arguments(tool_call) + + assert parsed is None + assert error == "arguments could not be parsed" + warning_messages = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert len(warning_messages) == 1 + assert warning_messages[0].startswith("Tool Permission Guardrail: Failed to decode arguments for tool Bash") + @pytest.mark.asyncio async def test_async_pre_call_hook_blocks_legacy_functions(self): data = { diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index acb45038df0..d9969dd1dc9 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5206,3 +5206,37 @@ class TestAzureRouterModelStreamingKeepalive: assert result.headers["x-upstream"] == "kept" assert chunks == [b"data: hello\n\n"] + + +@pytest.mark.asyncio +async def test_bedrock_count_tokens_error_forwards_provider_headers(): + """The count tokens route converts BedrockError into an HTTPException, and dropping the + headers there loses x-amzn-RequestId after the handler went to the trouble of keeping it.""" + from fastapi import HTTPException + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + failure = BedrockError( + status_code=500, + message="Amazon Bedrock is unable to process your request.", + headers={"x-amzn-RequestId": "req-count-tokens-500"}, + ) + + with patch( # test-quality-ok: the route's BedrockError branch is only reachable when the handler raises + "litellm.llms.bedrock.count_tokens.handler.BedrockCountTokensHandler.handle_count_tokens_request", + new=AsyncMock(side_effect=failure), + ): + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint="v1/messages/count_tokens", + request=MagicMock(), + fastapi_response=MagicMock(), + user_api_key_dict=MagicMock(), + request_body={"model": "anthropic.claude-haiku-4-5-20251001-v1:0"}, + ) + + assert exc_info.value.status_code == 500 + assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 54ef1f79f4d..61880a8c6f6 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -4,21 +4,27 @@ Tests for the pipeline executor. Uses mock guardrails to validate pipeline execution without external services. """ +import copy import logging +from typing import Literal from unittest.mock import MagicMock import pytest import litellm +from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( CustomCodeGuardrail, ) from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) +from litellm.types.utils import CallTypesLiteral try: from fastapi.exceptions import HTTPException @@ -159,11 +165,146 @@ class ContentCheckGuardrail(CustomGuardrail): return None +class RecordingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, scan_raw_request: bool = False, block: bool = True): + super().__init__( + guardrail_name=guardrail_name, + event_hook="pre_call", + default_on=True, + scan_raw_request=scan_raw_request, + ) + self.block = block + + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"detected": ["aws_access_key"]}, + request_data=data, + guardrail_status="guardrail_intervened" if self.block else "success", + ) + if self.block: + raise HTTPException(status_code=400, detail="Content policy violation") + return copy.deepcopy(data) + + # ───────────────────────────────────────────────────────────────────────────── # Tests # ───────────────────────────────────────────────────────────────────────────── +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +@pytest.mark.parametrize("scan_raw_request", [False, True]) +@pytest.mark.parametrize("on_fail", ["block", "modify_response"]) +async def test_terminal_block_carries_guardrail_information_to_request( + monkeypatch: pytest.MonkeyPatch, scan_raw_request: bool, on_fail: Literal["block", "modify_response"] +): + """ + Spend logging and the Guardrails Monitor read standard_logging_guardrail_information + off the caller's request dict. A blocking step records it on the executor's + working copy (or the raw-request snapshot), so the terminal result must carry it + back onto the request or the block is never counted. + """ + guard = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=scan_raw_request) + monkeypatch.setattr(litellm, "callbacks", [guard]) + data = { + "messages": [{"role": "user", "content": "key AKIAIOSFODNN7EXAMPLE"}], + "metadata": {"user_api_key_hash": "abc"}, + } + + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="credentials-api-keys", on_fail=on_fail, on_pass="next")], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={"messages": data["messages"], "metadata": {"user_api_key_hash": "abc"}}, + ) + + assert result.terminal_action == on_fail + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["credentials-api-keys"] + assert recorded[0]["guardrail_status"] == "guardrail_intervened" + assert data["metadata"]["user_api_key_hash"] == "abc" + assert "guardrails" not in data["metadata"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_terminal_block_merges_guardrail_information_without_duplicates(monkeypatch: pytest.MonkeyPatch): + """A pass_data step that returns a rewritten copy of the request, and a scan_raw_request step + that evaluates a deep copy taken before the pipeline ran, both leave earlier entries in two + dicts at once. Those must be carried back once while every step's own entry is kept.""" + first = RecordingGuardrail(guardrail_name="pii-scan", block=False) + second = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=True) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="pii-scan", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="credentials-api-keys", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "block" + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "pii-scan", "credentials-api-keys"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_repeated_scan_raw_request_step_is_counted_once_per_evaluation(monkeypatch: pytest.MonkeyPatch): + """Running the same raw-scan guardrail twice yields two identical entries; both must reach the caller, + while the entries the raw snapshot already held before the pipeline ran are not copied again.""" + guard = RecordingGuardrail(guardrail_name="credentials-raw", scan_raw_request=True, block=False) + monkeypatch.setattr(litellm, "callbacks", [guard]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="raw-scan-policy", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + recorded = result.modified_data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "credentials-raw", "credentials-raw"] + + @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio async def test_escalation_step1_fails_step2_blocks(monkeypatch): diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 2babfe432f3..ae4ec4086bb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -9,6 +9,7 @@ Pins covered: from __future__ import annotations import json +import logging import os import re from types import SimpleNamespace @@ -1633,6 +1634,32 @@ async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch): } +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", ["true", "false", "null", "'true'", None]) +async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monkeypatch, caplog, setting): + config_file = tmp_path / "budget.yaml" + flag = f" disable_budget_reservation: {setting}\n" if setting is not None else "" + config_file.write_text( + "model_list: []\nlitellm_settings: {}\ngeneral_settings:\n" + " master_key: null\n" + flag + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr("litellm.constants.budget_reservation_disabled_info_emitted", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + config = ProxyConfig() + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + for _ in range(3): + await config.load_config(router=None, config_file_path=str(config_file)) + + records = [ + record for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else []) + + @pytest.mark.asyncio async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch): """Regression: router_settings.plugins dotted-path strings must be resolved to @@ -2912,6 +2939,129 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing") +def _stub_add_deployment_collaborators( + monkeypatch: pytest.MonkeyPatch, pc: ProxyConfig, fake_prisma: MagicMock +) -> None: + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.get_model_list = MagicMock(return_value=[]) + + async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]: + return {} + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(pc, "get_config", fake_get_config) + monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock()) + monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock()) + monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "master_key", "sk-master") + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "proxy_config", pc) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + + +def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, object]: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + return { + "credential_name": credential_name, + "credential_values": {"api_key": encrypt_value_helper(api_key, new_encryption_key="sk-master")}, + "credential_info": {"custom_llm_provider": "openai"}, + } + + +def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock: + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row(credential_name, api_key)] + ) + return fake_prisma + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + from litellm.utils import load_credentials_from_list + + pc = ProxyConfig() + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {}) + installed = MagicMock() + + async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]: + fake_prisma.db.litellm_credentialstable.find_many.return_value = [ + _encrypted_credential_row("openai-cred", "sk-from-db") + ] + return [MagicMock()] + + async def install_models(new_models: object, proxy_logging_obj: object) -> None: + installed(credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands) + monkeypatch.setattr(pc, "_update_llm_router", install_models) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + installed.assert_called_once_with(credential={"api_key": "sk-from-db"}) + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + request_kwargs = {"litellm_credential_name": "openai-cred"} + load_credentials_from_list(request_kwargs) + assert request_kwargs == {"litellm_credential_name": "openai-cred", "api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_even_when_models_are_not_db_objects( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + + pc = ProxyConfig() + fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db") + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["mcp"]}) + models_fetch = AsyncMock(return_value=[]) + monkeypatch.setattr(pc, "_get_models_from_db", models_fetch) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + models_fetch.assert_not_awaited() + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_get_credentials_reads_from_writer_not_replica(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + pc = ProxyConfig() + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + writer_inner.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row("openai-cred", "sk-from-writer")] + ) + reader_inner.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + fake_prisma = MagicMock() + fake_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + + await pc.get_credentials(prisma_client=fake_prisma) + + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-writer"} + reader_inner.litellm_credentialstable.find_many.assert_not_awaited() + + # --------------------------------------------------------------------------- # ProxyConfig._add_general_settings_from_db_config # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index cc8fdeb0160..e466edab131 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -1162,6 +1162,104 @@ def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one(): assert result.prompt_caching > at_public_rates.prompt_caching +@pytest.mark.parametrize( + "baseline_id, selected_id, selected_multiplier, billed_input, classifier_cost, expected", + [ + ("baseline", "selected", 0.1, None, 0.0, 0.0135), + ("baseline", "selected", 2.0, None, 0.0, -0.015), + ("baseline", "selected", 1.0, None, 0.0, 0.0), + ("baseline", "selected", 0.1, 0.004, 0.001, 0.01), + ("baseline", "baseline", 0.1, 0.004, 0.001, -0.001), + (None, "selected", 0.1, None, 0.0, 0.0), + ("baseline", None, 0.1, None, 0.0, 0.0), + (None, None, 0.1, None, 0.0, 0.0), + ("", "selected", 0.1, None, 0.0, 0.0), + ("baseline", "", 0.1, None, 0.0, 0.0), + ], +) +def test_autorouter_savings_distinguishes_priced_deployments( + baseline_id: str | None, + selected_id: str | None, + selected_multiplier: float, + billed_input: float | None, + classifier_cost: float, + expected: float, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "anthropic/claude-opus-5", + "api_key": "test-key", + "input_cost_per_token": 1e-5 * multiplier, + "output_cost_per_token": 5e-5 * multiplier, + }, + "model_info": {"id": name}, + } + for name, multiplier in (("baseline", 1.0), ("selected", selected_multiplier)) + ] + ) + result: Final = compute_savings_spend( + model="claude-opus-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id=selected_id, + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "anthropic/claude-opus-5", + "savings_baseline_deployment_id": baseline_id, + "conversation_continuing": False, + "classifier_cost": classifier_cost, + }, + usage_object={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, + cost_breakdown=None if billed_input is None else {"input_cost": billed_input, "output_cost": 0.0}, + ) + assert result.autorouter == pytest.approx(expected) + + +@pytest.mark.parametrize("selected_model", ["azure/contract-deployment", "contract-deployment"]) +def test_autorouter_savings_recognizes_one_deployment_under_its_base_model(selected_model: str) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "contract", + "litellm_params": { + "model": "azure/contract-deployment", + "api_key": "test-key", + "api_base": "https://example.openai.azure.com", + "input_cost_per_token": 0.0001, + "output_cost_per_token": 0.0002, + "cache_read_input_token_cost": 0.00001, + }, + "model_info": {"id": "contract", "base_model": "azure/gpt-5.5"}, + } + ] + ) + result: Final = compute_savings_spend( + model=selected_model, + custom_llm_provider="azure", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id="contract", + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "azure/gpt-5.5", + "savings_baseline_deployment_id": "contract", + "conversation_continuing": True, + }, + usage_object={ + "prompt_tokens": 21000, + "completion_tokens": 100, + "total_tokens": 21100, + "prompt_tokens_details": {"text_tokens": 1000, "cached_tokens": 0, "cache_creation_tokens": 20000}, + }, + cost_breakdown={"input_cost": 2.1, "output_cost": 0.02}, + ) + assert result.autorouter == 0.0 + + def test_a_recorded_baseline_deployment_prices_at_its_configured_rate(): """A hardest-tier deployment with a negotiated rate is what the traffic would really have cost; pricing its model publicly misstates the saving.""" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 6acd9d7258e..bfae42f64f1 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8259,3 +8259,38 @@ class TestPassthroughHeadersAcceptImmutableMappings: assert merged["content-type"] == "text/event-stream" # the excluded hop-by-hop header is still dropped assert "transfer-encoding" not in merged + + +@pytest.mark.asyncio +async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status_error(): + """The httpx.HTTPStatusError branch dropped the headers its sibling branches forward. + + A Bedrock passthrough failure reaches this branch, so the request id was gone + before the client saw the response. + """ + import httpx + + from litellm.proxy._types import UserAPIKeyAuth + + request = httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse") + response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-passthrough-500"}, + content=b'{"message": "Amazon Bedrock is unable to process your request."}', + request=request, + ) + + processor = ProxyBaseLLMRequestProcessing(data={}) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(HTTPException) as exc_info: + await processor._handle_llm_api_exception( + e=httpx.HTTPStatusError("boom", request=request, response=response), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert exc_info.value.headers is not None + assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500" diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 9f1321aec2c..22930a26974 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == [] -def test_deployment_pre_call_target_stays_native_when_opted_out(): +def test_deployment_hook_target_stays_native_when_opted_out(): """Model-level guardrails resolve their target here rather than through ProxyLogging.""" - assert _KeepsNativeHooks()._deployment_pre_call_target() is not None + assert _KeepsNativeHooks()._deployment_hook_target() is not None opted_out = _KeepsNativeHooks() - assert opted_out._deployment_pre_call_target() is opted_out - assert _AppliesGuardrail()._deployment_pre_call_target() is not None + assert opted_out._deployment_hook_target() is opted_out + assert _AppliesGuardrail()._deployment_hook_target() is not None routed = _AppliesGuardrail() - assert routed._deployment_pre_call_target() is not routed + assert routed._deployment_hook_target() is not routed @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b7bb58378d4..937e4f15741 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -822,16 +822,16 @@ def _mock_scheduled_proxy_config() -> MagicMock: @pytest.mark.asyncio -async def test_initialize_scheduled_jobs_credentials(monkeypatch): - """ - Test that get_credentials is only called when store_model_in_db is True - """ +async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_deployment( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging - # Mock dependencies mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() @@ -841,25 +841,6 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch("litellm.proxy.proxy_server.store_model_in_db", False), - ): # set store_model_in_db to False - # Test when store_model_in_db is False - await ProxyStartupEvent.initialize_scheduled_background_jobs( - general_settings={}, - prisma_client=mock_prisma_client, - proxy_budget_rescheduler_min_time=1, - proxy_budget_rescheduler_max_time=2, - proxy_batch_write_at=5, - proxy_logging_obj=mock_proxy_logging, - ) - - # Verify get_credentials was not called - mock_proxy_config.get_credentials.assert_not_called() - - # Now test with store_model_in_db = True - with ( - patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), - patch("litellm.proxy.proxy_server.store_model_in_db", True), - patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True), ): await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={}, @@ -870,12 +851,31 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): proxy_logging_obj=mock_proxy_logging, ) - # Verify get_credentials was called both directly and scheduled - assert mock_proxy_config.get_credentials.call_count == 1 # Direct call + mock_proxy_config.get_credentials.assert_not_called() + mock_proxy_config.add_deployment.assert_not_called() - # Verify a scheduled job was added for get_credentials - mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls] - assert len(mock_scheduler_calls) > 0 + scheduler = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + assert scheduler.get_job("get_credentials_job") is None + assert scheduler.get_job("add_deployment_job") is not None + mock_proxy_config.get_credentials.assert_not_called() + assert mock_proxy_config.add_deployment.call_count == 1 + finally: + scheduler.shutdown(wait=False) @pytest.mark.asyncio @@ -924,7 +924,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat @pytest.mark.asyncio async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch): """ - The DB config-reload jobs (add_deployment, get_credentials) that keep multi-pod + The DB config-reload job (add_deployment) that keeps multi-pod deployments in sync must be scheduled at the configured proxy_config_reload_interval_seconds, not a hardcoded value. """ @@ -967,7 +967,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval( if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == configured_interval - assert scheduled_seconds["get_credentials_job"] == configured_interval + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -1011,7 +1011,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == 30 - assert scheduled_seconds["get_credentials_job"] == 30 + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -7446,10 +7446,8 @@ async def test_store_model_in_db_db_override_when_config_false(): # store_model_in_db should now be True (overridden by DB) assert ps.store_model_in_db is True - # add_deployment and get_credentials should have been called - # since store_model_in_db is now True assert mock_proxy_config.add_deployment.call_count == 1 - assert mock_proxy_config.get_credentials.call_count == 1 + mock_proxy_config.get_credentials.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index ab4c5185057..2c1845f7b92 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1,6 +1,13 @@ +import json +import sys +import types + import pytest +import respx +from httpx import Response from unittest.mock import AsyncMock, patch +import litellm from litellm.types.utils import ModelResponse from litellm.responses.mcp import chat_completions_handler @@ -1344,3 +1351,39 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti assert len(all_chunks) == 3 assert initial_stream.drained_after_exhaustion is True + + +@pytest.mark.asyncio +@respx.mock +async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + zapier_tool = {"type": "mcp", "server_label": "zapier", "server_url": "https://mcp.zapier.com/api/mcp/mcp"} + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-zapier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4.1", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + result = await acompletion_with_mcp( + model="openai/gpt-4.1", + messages=[{"role": "user", "content": "hello"}], + tools=[zapier_tool], + api_key="sk-test", + acompletion=True, + ) + + assert isinstance(result, ModelResponse) + assert result.id == "chatcmpl-zapier" + assert json.loads(provider.calls.last.request.content)["tools"] == [zapier_tool] diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 80151d0cba8..9745a0af970 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from openai.types.responses.tool_param import Mcp import importlib from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing @@ -724,6 +725,178 @@ def test_extract_tool_call_details_still_prefers_openai_arguments(): assert arguments == '{"city": "Paris"}' +def _registered( + server_id: str, + name: str, + alias: str | None = None, + server_name: str | None = None, + access_groups: list[str] | None = None, +): + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id=server_id, + name=name, + alias=alias, + server_name=server_name, + transport=MCPTransport.http, + access_groups=access_groups, + ) + + +async def _no_toolset(_: str) -> bool: + return False + + +ZAPIER_TOOL: Mcp = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "require_approval": "never", +} +EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"} +FUNCTION_TOOL = {"type": "function", "name": "get_weather", "parameters": {}} + + +@pytest.mark.asyncio +async def test_gateway_served_names_matches_alias_server_name_name_access_group_and_toolset(): + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + servers = ( + _registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]), + _registered("id-2", "deepwiki"), + ) + + async def toolset_exists(name: str) -> bool: + return name == "my-toolset" + + served = await _gateway_served_names( + {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"}, + servers=lambda: servers, + toolset_exists=toolset_exists, + ) + + assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"} + + +@pytest.mark.asyncio +async def test_gateway_served_names_matches_server_id_short_prefix_and_alias_case_like_the_gateway(): + from litellm.proxy._experimental.mcp_server.utils import compute_short_server_prefix + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b" + short_prefix = compute_short_server_prefix(server_id) + servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),) + + served = await _gateway_served_names( + {server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset + ) + + assert served == {server_id, short_prefix, "GitHub"} + + +@pytest.mark.asyncio +async def test_routes_through_gateway_flags_explicit_and_served_tools_only(): + served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"} + + async def served_names(names): + assert names == {"github", "mcp"} + return frozenset({"github"}) + + flags = await LiteLLM_Proxy_MCP_Handler.routes_through_gateway( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names + ) + + assert flags == (False, True, True, False) + + +@pytest.mark.asyncio +async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider(): + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway(): + served_tool = { + "type": "mcp", + "server_label": "toolset", + "server_url": "http://localhost:4000/mcp/my-toolset", + "require_approval": "never", + "allowed_tools": ["get_me"], + } + unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"} + + async def served_names(names): + return frozenset({"my-toolset"}) + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [served_tool, unserved_tool], served_names=served_names + ) + + assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}] + assert other_tools == [unserved_tool] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy(): + async def served_names(names): + raise AssertionError("no lookup expected") + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [FUNCTION_TOOL] + + +def test_should_use_gateway_still_triggers_on_http_mcp_path(): + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) is False + + +@pytest.mark.asyncio +async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.responses import main as responses_main + from litellm.types.llms.openai import ResponsesAPIResponse + + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + provider_tools: list[object] = [] + + def fake_provider(**kwargs: object) -> object: + request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"]) + provider_tools.append(request_params.get("tools")) + + async def respond() -> ResponsesAPIResponse: + return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[]) + + return respond() + + monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider) + + response = await responses_main.aresponses_api_with_mcp( + input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL] + ) + + assert isinstance(response, ResponsesAPIResponse) + assert provider_tools == [[ZAPIER_TOOL]] + + def _response_with_reasoning_and_tool_call() -> Any: """A first-turn response as a reasoning model returns it: reasoning item, then a function call.""" return ResponsesAPIResponse( diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index ebfb631f93b..5b1d8562abd 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -7,7 +7,8 @@ Tests the rule-based complexity scoring and tier assignment logic. import asyncio import logging import sys -from typing import Dict, List +import time +from typing import Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -47,6 +48,7 @@ from litellm.router_strategy.complexity_router.config import ( ClassifierLLMConfig, ComplexityRouterConfig, ComplexityTier, + custom_pattern_work, ) from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, @@ -756,6 +758,205 @@ class TestCustomTechnicalKeywords: assert custom_score > baseline_score +class TestCustomDimensions: + @pytest.mark.parametrize( + "matchers,prompt", + [ + pytest.param( + {"keywords": ["orbitmesh", "fluxgate"]}, + "Connect ORBITMESH and fluxgate for the requested change", + id="keywords", + ), + pytest.param( + {"patterns": [r"\bCREATE\s{1,4}TABLE\b", r"\bALTER\s{1,4}TABLE\b"]}, + "create table widgets (id integer); ALTER TABLE widgets ADD label text;", + id="regex", + ), + ], + ) + def test_custom_dimension_changes_only_matching_requests( + self, mock_router_instance: MagicMock, matchers: dict[str, object], prompt: str + ) -> None: + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + configured: Final = ComplexityRouter( + "test-router", + mock_router_instance, + {"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, **matchers}]}, + ) + baseline_tier, baseline_score, baseline_signals = baseline.classify(prompt) + tier, score, signals = configured.classify(prompt) + assert baseline_tier == ComplexityTier.SIMPLE + assert tier != ComplexityTier.SIMPLE + assert score == pytest.approx(baseline_score + 0.7) + assert signals == [*baseline_signals, "custom (internalFrameworks)"] + plain: Final = "Hello!" + assert configured.classify(plain) == baseline.classify(plain) + assert configured.classify(plain)[0] == ComplexityTier.SIMPLE + + @pytest.mark.parametrize( + "dimension_overrides,config_overrides", + [ + pytest.param({"keywords": []}, {}, id="missing-matchers"), + pytest.param({"keywords": [" "]}, {}, id="blank-keyword"), + pytest.param({"patterns": ["\t"]}, {}, id="blank-pattern"), + pytest.param({"patterns": ["("]}, {}, id="invalid-regex"), + pytest.param({"patterns": [r"a*b"]}, {}, id="unbounded-star"), + pytest.param({"patterns": [r"a{2,}b"]}, {}, id="unbounded-brace"), + pytest.param({"patterns": [r"a{0,65}b"]}, {}, id="repeat-over-64"), + pytest.param({"patterns": [r"(a{0,8}){0,8}b"]}, {}, id="nested-repeat"), + pytest.param({"patterns": [r"(a|aa){0,12}b"]}, {}, id="alternation-in-repeat"), + pytest.param({"patterns": [r"(?:ab){0,64}c"]}, {}, id="group-repeat"), + pytest.param({"patterns": ["a?" * 9 + "b"]}, {}, id="pattern-work-over-budget"), + pytest.param({"patterns": ["(?:a|aa)" * 9 + "z"]}, {}, id="ambiguous-alternation-chain"), + pytest.param({"patterns": ["a?" * 8 + "a{64}" * 10 + "z"]}, {}, id="cheap-prefix-expensive-tail"), + pytest.param({"patterns": [r"(a)\1"]}, {}, id="backreference"), + pytest.param({"patterns": [r"(?=x)y"]}, {}, id="lookahead"), + pytest.param({"patterns": [r"(?>ab)"]}, {}, id="atomic-group"), + pytest.param({"patterns": [r"a*+b"]}, {}, id="possessive"), + pytest.param({"name": "CODEPRESENCE"}, {"dimension_weights": {"tokenCount": 0.1}}, id="reserved-name"), + pytest.param({}, {"dimension_weights": {"INTERNALFRAMEWORKS": 0.7}}, id="weight-in-map"), + pytest.param({"weight": 0}, {}, id="zero-weight"), + pytest.param({"weight": 1.1}, {}, id="excess-weight"), + pytest.param({"weight": float("nan")}, {}, id="nan-weight"), + pytest.param({"weight": float("inf")}, {}, id="infinite-weight"), + pytest.param({"name": "bad-name"}, {}, id="invalid-name"), + pytest.param({"name": "x" * 65}, {}, id="long-name"), + pytest.param({"keywords": [""]}, {}, id="empty-matcher"), + pytest.param({"keywords": ["x" * 257]}, {}, id="long-matcher"), + pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"), + pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"), + pytest.param({"unknown": True}, {}, id="extra-field"), + ], + ) + def test_custom_dimension_invalid_configuration_rejected( + self, dimension_overrides: dict[str, object], config_overrides: dict[str, object] + ) -> None: + with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"): + ComplexityRouterConfig.model_validate( + { + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.7, + "keywords": ["orbitmesh"], + **dimension_overrides, + } + ], + **config_overrides, + } + ) + + @pytest.mark.parametrize( + "names", + [ + pytest.param(("internalFrameworks", "INTERNALFRAMEWORKS"), id="duplicate-casefolded-name"), + pytest.param(tuple(f"dimension{i}" for i in range(17)), id="dimension-count"), + ], + ) + def test_custom_dimension_names_and_count_are_bounded(self, names: tuple[str, ...]) -> None: + with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"): + ComplexityRouterConfig.model_validate( + {"custom_dimensions": [{"name": name, "weight": 0.7, "keywords": ["orbitmesh"]} for name in names]} + ) + + @pytest.mark.parametrize("classifier_type", ("heuristic_v2", "llm", "custom")) + def test_custom_dimensions_reject_classifiers_outside_the_tuning_gate(self, classifier_type: str) -> None: + classifier_config: Final = ( + {"classifier_plugin": _FixedTierClassifier("SIMPLE")} + if classifier_type == "custom" + else {"classifier_llm_config": {"model": "judge"}} + if classifier_type == "llm" + else {} + ) + with pytest.raises(ValidationError, match="custom_dimensions requires classifier_type"): + ComplexityRouterConfig.model_validate( + { + "classifier_type": classifier_type, + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], + **classifier_config, + } + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh")) + async def test_custom_dimensions_public_hook_scores_only_current_ask( + self, mock_router_instance: MagicMock, current_ask: str + ) -> None: + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], + }, + ) + result: Final = await router.async_pre_routing_hook( + model="test-router", + request_kwargs={}, + messages=[ + {"role": "system", "content": "orbitmesh"}, + {"role": "user", "content": "orbitmesh"}, + {"role": "assistant", "content": "orbitmesh is ready"}, + {"role": "user", "content": current_ask}, + {"role": "tool", "tool_call_id": "previous", "content": "orbitmesh"}, + ], + ) + assert result is not None + assert result.routing_decision is not None + assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (current_ask == "orbitmesh") + assert result.model == ("top" if current_ask == "orbitmesh" else "cheap") + assert "orbitmesh" not in " ".join(result.routing_decision["signals"]) + + def test_custom_patterns_scan_only_the_first_2048_characters(self, mock_router_instance: MagicMock) -> None: + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + {"custom_dimensions": [{"name": "late", "weight": 0.7, "patterns": [r"zzz{1,3}"]}]}, + ) + assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2] + assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2] + + def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None: + heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]} + ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(6)]}) + with pytest.raises(ValidationError, match="regex work estimate is 8939"): + ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(7)]}) + + @pytest.mark.parametrize( + "pattern,work", + [ + pytest.param(r"\b(create|alter|drop)\s{1,4}table\b", 135, id="sql-ddl"), + pytest.param("a?" * 8 + "z", 1277, id="optional-chain-near-cap"), + pytest.param(r"a{0,15}a{0,15}z", 801, id="adjacent-bounded-near-cap"), + pytest.param(r"[a-z0-9_]{3,63}\.(com|net|io)", 1291, id="class-repeat-plus-alternation"), + pytest.param("(?:a|aa)" * 8 + "z", 1787, id="ambiguous-alternation-near-cap"), + pytest.param("a{64}" * 10 + "z", 662, id="long-deterministic-tail"), + ], + ) + def test_custom_pattern_work_stays_cheap_on_adversarial_text( + self, mock_router_instance: MagicMock, pattern: str, work: int + ) -> None: + assert custom_pattern_work(pattern) == work + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "custom_dimensions": [ + {"name": "bounded", "weight": 0.7, "patterns": [pattern]}, + {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}, + ] + }, + ) + adversarial: Final = "orbitmesh " + "a" * 4000 + started: Final = time.perf_counter() + tier, score, signals = router.classify(adversarial) + elapsed: Final = time.perf_counter() - started + assert signals == ["long (1002 tokens)", "custom (internalFrameworks)"] + assert score == pytest.approx(0.8) + assert tier == ComplexityTier.REASONING + assert elapsed < 0.1 + + class TestAsyncPreRoutingHookEdgeCases: """Test edge cases for async_pre_routing_hook method.""" diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py new file mode 100644 index 00000000000..9efa526fc02 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -0,0 +1,187 @@ +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler + +GROUP: Final = "least-busy-group" +DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} +DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} +HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] + + +def _call_kwargs(deployment_id: str) -> dict[str, object]: + return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} + + +class SharedRedisCounters: + """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when + the key is created, and ordered reads that raise rather than invent a value.""" + + def __init__(self) -> None: + self.counts: dict[str, int] = {} + self.ttls: dict[str, int] = {} + + def count(self, key: str) -> int | None: + return self.counts.get(key) + + def expire(self, key: str) -> None: + self.counts.pop(key, None) + self.ttls.pop(key, None) + + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + incremented: Final = max(0, self.counts.get(key, 0) + value) + self.counts[key] = incremented + self.ttls.setdefault(key, ttl) + return incremented + + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + return self.increment_with_floor(key, value, ttl) + + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return tuple(self.counts.get(key) for key in key_list) + + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return self.batch_get_counts(key_list) + + +def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: + cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double + return LeastBusyLoggingHandler(router_cache=cache) + + +@pytest.mark.asyncio +async def test_worker_routes_around_a_request_another_worker_started() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + await picking_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await streaming_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_sync_pick_reads_the_shared_counts() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + picking_worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: + """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that + stuck count an hour after the key was created, which only works while nothing writes the TTL + again: a handler that refreshed it on every touch would keep the count alive for as long as + the group takes traffic, and the deployment would read busier than it is forever.""" + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + key: Final = f"{GROUP}_request_count:dep-a" + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + + shared.ttls[key] = 5 + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(key) == 1 + assert shared.ttls == {key: 5} + + +@pytest.mark.asyncio +async def test_counts_stay_in_memory_without_redis() -> None: + worker: Final = _worker(None) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + +class UnavailableRedis(SharedRedisCounters): + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + raise ConnectionError("redis is down") + + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + raise ConnectionError("redis is down") + + +@pytest.mark.asyncio +async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None: + worker: Final = _worker(UnavailableRedis()) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + shared.expire(f"{GROUP}_request_count:dep-a") + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 1 + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.asyncio +async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> None: + worker: Final = _worker(None) + in_memory: Final = worker.router_cache.in_memory_cache + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + in_memory.delete_cache(f"{GROUP}_request_count:dep-a") + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_calls_without_a_deployment_are_ignored() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) + worker.log_pre_api_call(model="m", messages=[], kwargs={}) + + assert shared.counts == {} diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index eb02459be68..812d7bbff32 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -165,6 +165,133 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): json.dumps({"latency": latencies}) +MODEL_GROUP = "gpt-4o-mini" +FAST_TTFT_ID = "fast-ttft-short-output" +SLOW_TTFT_ID = "slow-ttft-long-output" +STREAMING_DEPLOYMENTS = [ + {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}}, + {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}}, +] + + +def _streaming_kwargs(deployment_id: str, start_time: datetime, ttft_seconds: float): + return { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": deployment_id}, + }, + "stream": True, + "completion_start_time": start_time + timedelta(seconds=ttft_seconds), + } + + +def _recorded_ttft(cache: DualCache, deployment_id: str): + cached = cache.get_cache(key=f"{MODEL_GROUP}_map") or {} + return cached.get(deployment_id, {}).get("time_to_first_token_seconds", []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_streaming_ttft_ranking_ignores_completion_length(sync_mode: bool): + """Deployment A: TTFT 1s, 50 completion tokens. Deployment B: TTFT 3s, 500 + completion tokens. Dividing TTFT by completion tokens made B look faster + (3/500 = 0.006 beats 1/50 = 0.02); actual TTFT must win.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + start_time = datetime(2026, 1, 1, 12, 0, 0) + end_time = start_time + timedelta(seconds=10) + + samples = ( + (FAST_TTFT_ID, 1.0, 50), + (SLOW_TTFT_ID, 3.0, 500), + ) + for deployment_id, ttft, completion_tokens in samples: + kwargs = _streaming_kwargs(deployment_id, start_time, ttft) + response_obj = _chat_response(completion_tokens=completion_tokens) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(1.0)] + assert _recorded_ttft(cache, SLOW_TTFT_ID) == [pytest.approx(3.0)] + + request_kwargs = {"stream": True, "metadata": {}} + if sync_mode: + picked = handler.get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + else: + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_ttft_window_keeps_newest_samples_when_full(sync_mode: bool): + """Float timestamps, as the SDK passes them. Once max_latency_list_size + samples exist the oldest TTFT is dropped so the window slides.""" + max_size = 3 + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": max_size}) + start_time = 1_700_000_000.0 + ttfts = (0.1, 0.2, 0.3, 0.4) + + for ttft in ttfts: + kwargs = { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": FAST_TTFT_ID}, + }, + "stream": True, + "completion_start_time": start_time + ttft, + } + response_obj = _chat_response(completion_tokens=1) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(ttft) for ttft in ttfts[-max_size:]] + + +@pytest.mark.asyncio +async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_workers(): + """Workers on the previous release share the Redis map and keep writing + seconds-per-token under the old "time_to_first_token" key during a rolling + deploy. Those samples favor SLOW; routing must only read the seconds key.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token": [0.02], "time_to_first_token_seconds": [1.0]}, + SLOW_TTFT_ID: {"time_to_first_token": [0.006], "time_to_first_token_seconds": [3.0]}, + }, + ) + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + @pytest.mark.asyncio @pytest.mark.parametrize( "cached_entry", diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 5599c5aad63..af390c3292b 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy @@ -435,6 +436,81 @@ def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeyp assert router._group_selectors == {} +def test_two_least_busy_groups_count_a_request_once(monkeypatch): + """ + Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and + back down from the success hooks on `litellm.callbacks`. The success list drops a second + selector of the same class, so a pre-call list that kept both counted every request twice + and released it once, and the deployment's in-flight count climbed until it looked pinned. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + router = _build_router( + routing_strategy="least-busy", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "least-busy", + } + ], + ) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + +def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): + """ + Least-busy hangs its counting off litellm's global callback lists, and those lists keep one + logger per class unless the instances differ in a plain attribute. Two routers in one process + (a second Router, or a per-request `user_config` one) therefore have to register separately: + a second router whose selector is dropped counts nothing, reads zero for every deployment, + and sends every request to whichever one is listed first. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + first = _build_router(routing_strategy="least-busy") + second = _build_router(routing_strategy="least-busy") + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 1 + assert ( + second.get_available_deployment(model="filtered-model", messages=[])["model_info"]["id"] + == "deploy-2" + ) + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert first.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index 293af36080a..0a54addce5e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -56,6 +56,18 @@ class BlockEverything: return context +class MessageRecorder: + """Records what each plugin pass was handed, then blocks so the request stops there.""" + + def __init__(self): + self.seen = [] + + async def run(self, context: RoutingContext) -> RoutingContext: + self.seen.append(list(context.raw_messages)) + context.candidate_models = [] + return context + + def _smart_router_model_list(): return [ { @@ -164,6 +176,71 @@ async def test_async_completion_with_unsupported_strategy_rejects_configured_plu await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) +@pytest.mark.asyncio +async def test_prompt_management_model_still_runs_the_plugin_pipeline(): + """ + A prompt-management model routes through its own factory, which picked the deployment + on the synchronous path. Plugins never run there, so the guard turned every such request + into an error message about the caller's own API choice, on an async call the caller made + correctly. It also read the in-flight counts with a blocking call inside the event loop. + """ + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[BlockEverything()], + ) + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="lit-7039", + ) + + +@pytest.mark.asyncio +async def test_prompt_management_plugins_see_the_callers_own_messages(): + """ + The prompt-management factory picks its deployment with a placeholder message, which was + harmless while that pick ran on the synchronous path (plugins never ran there at all). Now + that the pick runs the plugin pipeline, a plugin that classifies request content would score + the placeholder instead of the conversation, and the narrowing it produces decides which + deployments the real call is allowed to use. + """ + recorder = MessageRecorder() + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[recorder], + ) + messages = [{"role": "user", "content": "wire me $40,000 to account 12345"}] + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=messages, + litellm_call_id="lit-7039", + ) + + assert recorder.seen == [messages] + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 79ae00e155c..030bdfe03e9 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -332,6 +332,67 @@ async def test_per_request_enable_prompt_caching_reaches_the_affinity_key(monkey assert filtered == [deployments[1]] +@pytest.mark.asyncio +async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_affinity_key(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) + request_kwargs = { + "system": [ + { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=true;", + } + ], + "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.263 (external, cli)"}}, + } + auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=(AUTO_CACHING_MODEL,), + ) + assert auto_injected_messages != messages + await PromptCachingCache(cache=cache).async_add_model_id( + model_id="dep-2", messages=auto_injected_messages, tools=None + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs=request_kwargs, + ) + + assert filtered == deployments + + +@pytest.mark.asyncio +async def test_root_cache_control_does_not_reuse_an_auto_injected_affinity_key(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + messages = _auto_caching_messages() + auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=(AUTO_CACHING_MODEL,), + ) + assert auto_injected_messages != messages + await PromptCachingCache(cache=cache).async_add_model_id( + model_id="dep-2", messages=auto_injected_messages, tools=None + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs={"cache_control": {"type": "ephemeral"}}, + ) + + assert filtered == deployments + + @pytest.mark.asyncio async def test_tool_marked_cache_control_keeps_routing_off_another_requests_prefix(monkeypatch, local_model_cost_map): """ diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py index fa7a96adb20..f686a62db76 100644 --- a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py +++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Mapping +from typing import Final import pytest @@ -58,6 +59,7 @@ class TestTuningFingerprint: "reasoning_override_min_score": 0.05, "token_thresholds": {"simple": 20, "complex": 500}, "dimension_weights": {"codePresence": 0.9}, + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], "code_keywords": ["orionflow"], "reasoning_keywords": ["deduce"], "technical_keywords": ["ledgerkit"], @@ -216,6 +218,28 @@ class TestQuota: is None ) + def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self) -> None: + baselines: Final = snapshot_tuning_baselines(()) + original: Final = _router("a", {}) + config: Final = { + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}] + } + edited_config: Final = { + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.9, "keywords": ["orbitmesh"]}] + } + added: Final = _router("a", config) + edited: Final = _router("a", edited_config) + second: Final = _router("b", config) + + assert tuning_fingerprint(config) != tuning_fingerprint(edited_config) + assert mutable_tuned_identities((added,), baselines) == {router_identity(original)} + assert tuning_quota_violation(candidate=added, others=(original,), baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited, others=(added,), baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=second, others=(edited,), baselines=baselines, limit=1) is not None + assert tuning_quota_violation(candidate=original, others=(edited,), baselines=baselines, limit=1) is None + assert mutable_tuned_identities((original,), baselines) == frozenset() + assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None + def test_violation_message_names_the_limit_and_remedy(self) -> None: message = tuning_limit_violation(held=2, limit=1) assert message is not None diff --git a/tests/test_litellm/router_utils/test_cooldown_cache.py b/tests/test_litellm/router_utils/test_cooldown_cache.py index 68e9aeaa4fc..6f90fa8465f 100644 --- a/tests/test_litellm/router_utils/test_cooldown_cache.py +++ b/tests/test_litellm/router_utils/test_cooldown_cache.py @@ -268,12 +268,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -289,7 +289,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -312,14 +312,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -340,12 +340,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None @pytest.mark.asyncio async def test_async_active_entry_is_returned(self): @@ -363,7 +363,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -389,18 +389,18 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:expired-dep:cooldown" entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is None - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None def test_active_entry_within_window_returns_value(self): cc = self._make_cooldown_cache() key = "deployment:active-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) @@ -412,12 +412,12 @@ class TestCorrectedActiveCooldown: key = "deployment:backfilled-dep:cooldown" remaining = 30.0 entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is not None - corrected_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + corrected_expiry = cc.in_memory_cache.ttl_dict.get(key) assert corrected_expiry is not None assert corrected_expiry - time.time() <= 60.0 @@ -425,10 +425,160 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:normal-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) - original_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) + original_expiry = cc.in_memory_cache.ttl_dict.get(key) cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry == original_expiry + + +class SharedRedisDouble: + """ + In-process stand-in for RedisCache, shared by several DualCache instances so that + tests can model two proxy replicas talking to one Redis. + """ + + def __init__(self) -> None: + self.store: dict = {} # mutable-ok: stands in for Redis' own mutable keyspace + + def set_cache(self, key, value, **kwargs): + self.store[key] = value + + async def async_set_cache(self, key, value, **kwargs): + self.store[key] = value + + def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + +class TestCooldownPropagationBetweenReplicas: + """ + A cooldown written by one replica has to reach its siblings quickly. The router's own + DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads + get their own cache with a much shorter Redis read interval. + """ + + def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache: + router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis) + if read_interval is None: + return CooldownCache(cache=router_cache, default_cooldown_time=60.0) + return CooldownCache( + cache=router_cache, + default_cooldown_time=60.0, + redis_read_interval_seconds=read_interval, + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "shared-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "sibling replica must pick up a cooldown written by another replica within the read interval" + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_default_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis) + replica_b = self._make_replica(redis) + model_id = "default-interval-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(1.2) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "the shipped default read interval must let a sibling replica see a cooldown about a second later" + ) + + def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "sync-shared-deployment" + + assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active] + + @pytest.mark.asyncio + async def test_redis_attached_after_construction_is_still_used(self): + redis = SharedRedisDouble() + router_cache = DualCache(in_memory_cache=InMemoryCache()) + writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25) + router_cache.attach_redis_cache(redis) + reader = self._make_replica(redis, read_interval=0.25) + model_id = "late-redis-deployment" + + writer.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "a router that wires Redis after building its cooldown cache must still publish cooldowns to it" + ) + + +class TestCooldownSurvivesUnrelatedCacheTraffic: + @pytest.mark.asyncio + async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self): + router_cache = DualCache(in_memory_cache=InMemoryCache()) + cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0) + model_id = "busy-router-deployment" + + cc.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=30.0, + ) + + for i in range(400): + router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i}) + + active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "unrelated router cache traffic must not evict a cooldown that is still running" + ) diff --git a/tests/test_litellm/test_bedrock_extended_beta_models.py b/tests/test_litellm/test_bedrock_extended_beta_models.py deleted file mode 100644 index ebbbd6cab5c..00000000000 --- a/tests/test_litellm/test_bedrock_extended_beta_models.py +++ /dev/null @@ -1,170 +0,0 @@ -""" -Test suite for AWS Bedrock extended beta model support -Tests model configuration, pricing, and regional availability for: -- DeepSeek V3.2 -- Minimax M2.1 -- Moonshot AI Kimi K2.5 -- Qwen3 Coder Next -""" - -import os - -# Set env var to use local model cost map instead of fetching from remote -os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true" - -import pytest - -from litellm import get_model_info - -# Model configurations: (model_name, regions, max_input, max_output) -MODEL_CONFIGS = [ - ( - "deepseek.v3.2", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-north-1", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 163840, - 163840, - ), - ( - "minimax.minimax-m2.1", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-central-1", - "eu-north-1", - "eu-south-1", - "eu-west-1", - "eu-west-2", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 196000, - 8192, - ), - ( - "moonshotai.kimi-k2.5", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-north-1", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 262144, - 262144, - ), - ( - "qwen.qwen3-coder-next", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-central-1", - "eu-south-1", - "eu-west-1", - "eu-west-2", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 262144, - 8192, - ), -] - - -class TestBedrockNewModels: - """Unified test suite for all new Bedrock models""" - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_model_info_primary_region( - self, model_name, regions, max_input, max_output - ): - """Test model configuration in primary region (us-east-1)""" - model = f"bedrock/us-east-1/{model_name}" - model_info = get_model_info(model) - - assert model_info is not None, f"Model {model_name} not found" - assert model_info["max_input_tokens"] == max_input - assert model_info["max_output_tokens"] == max_output - assert model_info["litellm_provider"] == "bedrock" - assert model_info["mode"] == "chat" - assert model_info["supports_function_calling"] is True - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_pricing_configured(self, model_name, regions, max_input, max_output): - """Verify pricing is set for all models""" - model = f"bedrock/us-east-1/{model_name}" - model_info = get_model_info(model) - - assert ( - model_info["input_cost_per_token"] > 0 - ), f"Missing input cost for {model_name}" - assert ( - model_info["output_cost_per_token"] > 0 - ), f"Missing output cost for {model_name}" - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_region_count(self, model_name, regions, max_input, max_output): - """Verify each bedrock/{region}/{model_name} resolves via get_model_info""" - for region in regions: - model = f"bedrock/{region}/{model_name}" - model_info = get_model_info(model) - assert model_info is not None, f"Model {model_name} not found in {region}" - assert model_info["max_input_tokens"] == max_input - assert model_info["max_output_tokens"] == max_output - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_sample_regional_variants(self, model_name, regions, max_input, max_output): - """Test sample regional variants (us-east-1, eu-west-1, ap-northeast-1)""" - for region in ["us-east-1", "ap-northeast-1"]: - if region in regions: - model = f"bedrock/{region}/{model_name}" - model_info = get_model_info(model) - assert ( - model_info is not None - ), f"Model {model_name} not found in {region}" - assert model_info["max_input_tokens"] == max_input - assert model_info["litellm_provider"] == "bedrock" - - -class TestModelSpecificFeatures: - """Model-specific capability tests""" - - def test_deepseek_v3_2_context_window(self): - """DeepSeek V3.2 has 163K context window""" - model_info = get_model_info("bedrock/us-east-1/deepseek.v3.2") - assert model_info["max_input_tokens"] == 163840 - - def test_minimax_m2_1_context_window(self): - """Minimax M2.1 has 196K input, 8K output""" - model_info = get_model_info("bedrock/us-east-1/minimax.minimax-m2.1") - assert model_info["max_input_tokens"] == 196000 - assert model_info["max_output_tokens"] == 8192 - - def test_moonshotai_kimi_k2_5_context_window(self): - """Moonshot AI Kimi K2.5 has 256K context window""" - model_info = get_model_info("bedrock/us-east-1/moonshotai.kimi-k2.5") - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - - def test_qwen3_coder_next_context_window(self): - """Qwen3 Coder Next has 256K input, 8K output""" - model_info = get_model_info("bedrock/us-east-1/qwen.qwen3-coder-next") - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 8192 diff --git a/tests/test_litellm/test_bedrock_nemotron_super.py b/tests/test_litellm/test_bedrock_nemotron_super.py deleted file mode 100644 index 969db890e84..00000000000 --- a/tests/test_litellm/test_bedrock_nemotron_super.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -Test suite for NVIDIA Nemotron Super 3 120B on AWS Bedrock -Verifies model configuration, pricing, and regional availability. -""" - -import os - -os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true" - -import pytest - -from litellm import get_model_info - - -MODEL_NAME = "nvidia.nemotron-super-3-120b" - - -class TestNemotronSuper3120B: - """Test model definition for nvidia.nemotron-super-3-120b""" - - def test_model_info_primary_region(self): - """Test model resolves in us-east-1""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info is not None, f"Model {MODEL_NAME} not found" - assert model_info["max_input_tokens"] == 256000 - assert model_info["max_output_tokens"] == 32768 - assert model_info["litellm_provider"] == "bedrock_converse" - assert model_info["mode"] == "chat" - assert model_info["supports_function_calling"] is True - - def test_pricing_configured(self): - """Verify pricing matches AWS Bedrock rates""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info["input_cost_per_token"] == 1.5e-07 - assert model_info["output_cost_per_token"] == 6.5e-07 - - def test_context_window(self): - """Nemotron Super 3 120B has 256K input, 32K output on Bedrock""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info["max_input_tokens"] == 256000 - assert model_info["max_output_tokens"] == 32768 - - def test_resolves_without_region(self): - """Test model resolves with just bedrock/ prefix""" - model_info = get_model_info(f"bedrock/{MODEL_NAME}") - - assert model_info is not None, f"Model {MODEL_NAME} not found without region" - assert model_info["max_input_tokens"] == 256000 diff --git a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py b/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py deleted file mode 100644 index 1312aa110d3..00000000000 --- a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Validate that AWS GovCloud (Bedrock us-gov-*) Haiku 4.5 entries carry -the 1-hour cache write tier. - -AWS Bedrock GovCloud pricing applies a +20% premium over global -Anthropic rates. Global Haiku 4.5 1h cache write is $2.00/MTok; us-gov -is therefore $2.40/MTok — exactly 1.6x the 5-minute rate of $1.50/MTok. - -Source: https://aws.amazon.com/bedrock/pricing/ -""" - -import json -import os - -import pytest - - -@pytest.fixture(scope="module") -def model_data(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - return json.load(f) - - -HAIKU_USGOV_KEYS = [ - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0", -] - - -@pytest.mark.parametrize("model_key", HAIKU_USGOV_KEYS) -def test_usgov_haiku_4_5_1hr_cache_write(model_data, model_key): - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - assert ( - info["cache_creation_input_token_cost"] == 1.5e-06 - ), f"{model_key}: 5m cache write should be $1.50/MTok" - assert ( - info["cache_creation_input_token_cost_above_1hr"] == 2.4e-06 - ), f"{model_key}: 1h cache write should be $2.40/MTok" - ratio = ( - info["cache_creation_input_token_cost_above_1hr"] - / info["cache_creation_input_token_cost"] - ) - assert abs(ratio - 1.6) < 1e-9, f"{model_key}: 1h/5m ratio is {ratio}, expected 1.6" diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py index 3576834dd27..3dfd7350a06 100644 --- a/tests/test_litellm/test_bedrock_usgov_pricing.py +++ b/tests/test_litellm/test_bedrock_usgov_pricing.py @@ -31,34 +31,6 @@ def model_data(): return json.load(f) -SONNET_4_5_USGOV_KEYS = [ - "bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0", - "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0", -] - - -@pytest.mark.parametrize("model_key", SONNET_4_5_USGOV_KEYS) -def test_usgov_sonnet_4_5_pricing(model_data, model_key): - """Each us-gov sonnet-4-5 entry must carry the +20%-over-global rates - that AWS publishes on the GovCloud pricing page. - """ - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - - assert info["input_cost_per_token"] == 3.6e-06, ( - f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})" - ) - assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok" - assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok" - assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, ( - f"{model_key}: 1h cache write should be $7.20/MTok" - ) - assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok" - - def test_usgov_carries_20_percent_premium_over_global(model_data): """The us-gov rates must equal 1.2x the global anthropic.* rates, matching AWS's documented GovCloud uplift. @@ -92,18 +64,6 @@ EXPECTED_USGOV_ABOVE_200K = { } -@pytest.mark.parametrize("field,expected", EXPECTED_USGOV_ABOVE_200K.items()) -def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, expected): - """The `_above_200k_tokens` tier on the us-gov cross-region inference - profile must also carry the +20% GovCloud uplift. The original PR - corrected the base rates but left the 200k-tier fields at the +10% - commercial-US rates, undercharging long-context requests. - """ - info = model_data[USGOV_CROSS_REGION_KEY] - assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}" - assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})" - - def test_usgov_cross_region_above_200k_ratio_to_global(model_data): """Cross-check via the property-based invariant: every `_above_200k_tokens` field on the us-gov cross-region profile must equal 1.2x the global @@ -117,167 +77,6 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data): assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" -CLAUDE_GOV_EXPECTED = { - "anthropic.claude-sonnet-5": { - "input_cost_per_token": 2.4e-06, - "output_cost_per_token": 1.2e-05, - "cache_creation_input_token_cost": 3e-06, - "cache_creation_input_token_cost_above_1hr": 4.8e-06, - "cache_read_input_token_cost": 2.4e-07, - }, - "anthropic.claude-opus-4-8": { - "input_cost_per_token": 6e-06, - "output_cost_per_token": 3e-05, - "cache_creation_input_token_cost": 7.5e-06, - "cache_creation_input_token_cost_above_1hr": 1.2e-05, - "cache_read_input_token_cost": 6e-07, - }, - "anthropic.claude-opus-5": { - "input_cost_per_token": 6e-06, - "output_cost_per_token": 3e-05, - "cache_creation_input_token_cost": 7.5e-06, - "cache_creation_input_token_cost_above_1hr": 1.2e-05, - "cache_read_input_token_cost": 6e-07, - }, - "anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.2e-05, - "output_cost_per_token": 6e-05, - "cache_creation_input_token_cost": 1.5e-05, - "cache_creation_input_token_cost_above_1hr": 2.4e-05, - "cache_read_input_token_cost": 3e-07, - }, -} - - -USGOV_CLAUDE_KEY_TEMPLATES = { - "bedrock/us-gov-east-1/{base_key}": "bedrock", - "bedrock/us-gov-west-1/{base_key}": "bedrock", - "us-gov.{base_key}": "bedrock_converse", -} - - -@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED) -@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) -def test_usgov_claude_pricing(model_data, key_template, expected_provider, base_key): - """Sonnet 5, Opus 4.8, Opus 5, and Fable 5.1 gov entries, both in-region keys - and the us-gov. geo inference profile the model cards list for GovCloud, must - carry the 1.2x GovCloud premium over the global anthropic.* rates. No public - AWS source (offer files, pricing page) lists Claude GovCloud rows; the premium - is the one AWS quotes for Opus 4.8 in GovCloud ($6/$30 per million). - """ - gov_key = key_template.format(base_key=base_key) - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == expected_provider - assert "search_context_cost_per_query" not in info - for field, expected in CLAUDE_GOV_EXPECTED[base_key].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - ratio = info[field] / model_data[base_key][field] - assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2" - - -CONVERSE_GOV_EXPECTED = { - "nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07), - "nvidia.nemotron-nano-9b-v2": (7.2e-08, 2.76e-07), - "nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07), - "nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07), - "openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07), - "openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07), -} - - -@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED) -@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) -def test_usgov_converse_model_pricing(model_data, key_template, expected_provider, base_key): - """Nemotron and gpt-oss gov entries, in-region and the us-gov. geo inference - profile both GovCloud regions list as ACTIVE, must match the AWS Bedrock - offer file, which prices both regions identically at 1.2x commercial. - """ - gov_key = key_template.format(base_key=base_key) - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key] - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["litellm_provider"] == expected_provider - base = model_data[base_key] - assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9 - assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9 - - -def test_usgov_west_llama3_8b_output_price_fixed(model_data): - """The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok); - the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model - in us-gov-west-1 only, so there is no east entry to check. - """ - info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"] - assert info["input_cost_per_token"] == 3e-07 - assert info["output_cost_per_token"] == 6e-07 - - -MANTLE_GOV_TIERED_EXPECTED = { - "openai.gpt-5.6-luna": { - "input_cost_per_token": 2.64e-07, - "input_cost_per_token_above_272k_tokens": 5.28e-07, - "cache_creation_input_token_cost": 3.3e-07, - "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, - "cache_read_input_token_cost": 2.64e-08, - "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, - "output_cost_per_token": 1.584e-06, - "output_cost_per_token_above_272k_tokens": 2.376e-06, - }, - "openai.gpt-5.6-terra": { - "input_cost_per_token": 2.64e-06, - "input_cost_per_token_above_272k_tokens": 5.28e-06, - "cache_creation_input_token_cost": 3.3e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, - "cache_read_input_token_cost": 2.64e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, - "output_cost_per_token": 1.584e-05, - "output_cost_per_token_above_272k_tokens": 2.376e-05, - }, -} - - -@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED) -def test_usgov_west_mantle_terra_luna_pricing(model_data, model): - """Terra and Luna carry 1.2x commercial across every tier in the - us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them. - """ - gov_key = f"bedrock_mantle/us-gov-west-1/{model}" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - assert info["litellm_provider"] == "bedrock_mantle" - assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data - - -@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) -def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region): - """gpt-5.4 gov rates come from the offer file, which publishes only the - standard tier in GovCloud: no long-context SKUs exist there, unlike commercial. - """ - gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["input_cost_per_token"] == 3.3e-06 - assert info["cache_read_input_token_cost"] == 3.3e-07 - assert info["output_cost_per_token"] == 1.98e-05 - assert not any(field.endswith("_above_272k_tokens") for field in info) - - -def test_usgov_mantle_grok_4_3_west_only(model_data): - """grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer - file carries grok-4.6 instead. - """ - info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 3e-06 - assert info["cache_read_input_token_cost"] == 2.4e-07 - assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data - - def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data): """us-gov-east-1 serves claude-3-haiku through the us-gov. inference profile only, so the profile row must bill exactly like the in-region gov row. @@ -290,96 +89,6 @@ def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data): } -GROK_4_6_GOV_KEYS = { - "us-gov.xai.grok-4.6": ("us.xai.grok-4.6", "bedrock_converse"), - "bedrock_mantle/us-gov-west-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"), - "bedrock_mantle/us-gov-east-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"), -} - - -@pytest.mark.parametrize("gov_key", GROK_4_6_GOV_KEYS) -def test_usgov_grok_4_6_pricing(model_data, gov_key): - """Both GovCloud regions serve grok-4.6 through the us-gov. profile only, and - both offer files price its standard SKU at 1.2x the commercial US rate. - """ - base_key, expected_provider = GROK_4_6_GOV_KEYS[gov_key] - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == expected_provider - assert info["input_cost_per_token"] == 2.64e-06 - assert info["output_cost_per_token"] == 7.92e-06 - assert info["cache_read_input_token_cost"] == 6.6e-07 - for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): - assert abs(info[field] / model_data[base_key][field] - 1.2) < 1e-9 - - -NOVA_GOV_WEST_EXPECTED = { - "amazon.nova-lite-v1:0": (7.2e-08, 2.88e-07), - "amazon.nova-micro-v1:0": (4.2e-08, 1.68e-07), -} - - -@pytest.mark.parametrize("base_key", NOVA_GOV_WEST_EXPECTED) -def test_usgov_west_nova_lite_micro_pricing(model_data, base_key): - """Nova Lite and Micro are on-demand in us-gov-west-1 only; the offer file - prices them at 1.2x commercial, like the Nova Pro row that was already there. - """ - gov_key = f"bedrock/us-gov-west-1/{base_key}" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - expected_input, expected_output = NOVA_GOV_WEST_EXPECTED[base_key] - assert info["litellm_provider"] == "bedrock" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert abs(info["input_cost_per_token"] / model_data[base_key]["input_cost_per_token"] - 1.2) < 1e-9 - assert abs(info["output_cost_per_token"] / model_data[base_key]["output_cost_per_token"] - 1.2) < 1e-9 - assert f"bedrock/us-gov-east-1/{base_key}" not in model_data - - -def test_usgov_west_nova_2_multimodal_embeddings_pricing(model_data): - """Every meter of the multimodal embedding model (tokens, images, audio and - video seconds) carries the 1.2x uplift the us-gov-west-1 offer file lists. - """ - gov_key = "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == "bedrock" - assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == 1.62e-07 - assert info["input_cost_per_image"] == 7.2e-05 - assert info["input_cost_per_audio_per_second"] == 0.000168 - assert info["input_cost_per_video_per_second"] == 0.00084 - assert "bedrock/us-gov-east-1/amazon.nova-2-multimodal-embeddings-v1:0" not in model_data - - -MANTLE_GOV_FLAT_EXPECTED = { - "google.gemma-4-e2b": (4.8e-08, 9.6e-08, ("us-gov-west-1",)), - "google.gemma-4-26b-a4b": (1.56e-07, 4.8e-07, ("us-gov-west-1",)), - "google.gemma-4-31b": (1.68e-07, 4.8e-07, ("us-gov-west-1",)), - "openai.gpt-oss-20b": (8.4e-08, 3.6e-07, ("us-gov-west-1", "us-gov-east-1")), - "openai.gpt-oss-120b": (1.8e-07, 7.2e-07, ("us-gov-west-1", "us-gov-east-1")), -} - - -@pytest.mark.parametrize("model", MANTLE_GOV_FLAT_EXPECTED) -def test_usgov_mantle_gemma_and_gpt_oss_pricing(model_data, model): - """Gemma 4 is priced in the us-gov-west-1 offer file only and gpt-oss in both; - each Mantle gov row carries the offer file's standard SKU, and no row exists - for a region whose offer file has no SKU. - """ - expected_input, expected_output, regions = MANTLE_GOV_FLAT_EXPECTED[model] - for region in ("us-gov-west-1", "us-gov-east-1"): - gov_key = f"bedrock_mantle/{region}/{model}" - if region not in regions: - assert gov_key not in model_data - continue - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == "bedrock_mantle" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - - GOV_ROW_SOURCES = { "us-gov.anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1", "bedrock/us-gov-west-1/anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1", @@ -417,33 +126,3 @@ def test_usgov_rows_keep_commercial_limits_and_capabilities(model_data, gov_key) assert _non_pricing_fields(gov) == _non_pricing_fields(model_data[GOV_ROW_SOURCES[gov_key]]) assert "search_context_cost_per_query" not in gov assert "source" not in gov - - -AZURE_GOV_EXPECTED = { - "azure/us-gov/gpt-5.1": { - "input_cost_per_token": 1.71875e-06, - "cache_read_input_token_cost": 1.71875e-07, - "output_cost_per_token": 1.375e-05, - }, - "azure/us-gov/o3-mini": { - "input_cost_per_token": 1.513e-06, - "cache_read_input_token_cost": 7.57e-07, - "output_cost_per_token": 6.05e-06, - }, - "azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07}, - "azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08}, -} - - -@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED) -def test_azure_usgov_pricing(model_data, gov_key): - """Azure Government meters from the Azure retail prices API - (usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government - retirement schedule is published, so these entries carry no deprecation_date. - """ - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - for field, expected in AZURE_GOV_EXPECTED[gov_key].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - assert info["litellm_provider"] == "azure" - assert "deprecation_date" not in info diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py index 3ecf94602d9..0473161faac 100644 --- a/tests/test_litellm/test_claude_fable_5_config.py +++ b/tests/test_litellm/test_claude_fable_5_config.py @@ -14,7 +14,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -27,89 +26,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_fable_5_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = [ - ("claude-fable-5", "anthropic"), - ("anthropic.claude-fable-5", "bedrock_converse"), - ("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"), - # Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context - # window on Microsoft Foundry. - ("azure_ai/claude-fable-5", "azure_ai"), - ] - - for model_name, provider in expected_models: - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m - # cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 1e-05 - assert info["output_cost_per_token"] == 5e-05 - assert info["cache_creation_input_token_cost"] == 1.25e-05 - assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 - assert info["cache_read_input_token_cost"] == 1e-06 - - # Flat-rate across the full 1M context window. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - - -def test_fable_5_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - # Fable 5 launched with us/eu geo inference profiles plus a global profile - # (no au/apac/jp). Global uses base pricing; geo profiles carry the - # standard 10% regional premium. - expected_models = { - "global.anthropic.claude-fable-5": { - "input_cost_per_token": 1e-05, - "output_cost_per_token": 5e-05, - "cache_creation_input_token_cost": 1.25e-05, - "cache_read_input_token_cost": 1e-06, - }, - "us.anthropic.claude-fable-5": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 1.1e-06, - }, - "eu.anthropic.claude-fable-5": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 1.1e-06, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - def test_fable_5_geo_multiplier_without_fast_mode(): """First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key @@ -144,13 +60,6 @@ def test_fable_5_registered_for_bedrock_converse(): assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS -def test_fable_5_provider_resolves_via_model_info(local_model_cost_map): - info = litellm.get_model_info(model="claude-fable-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], @@ -222,46 +131,6 @@ FABLE_5_1_VARIANTS = ( ) -def test_fable_5_1_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = [ - ("claude-fable-5-1", "anthropic"), - ("anthropic.claude-fable-5-1", "bedrock_converse"), - ("vertex_ai/claude-fable-5-1", "vertex_ai-anthropic_models"), - ("azure_ai/claude-fable-5-1", "azure_ai"), - ] - - for model_name, provider in expected_models: - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["input_cost_per_token"] == 1e-05 - assert info["output_cost_per_token"] == 5e-05 - assert info["cache_creation_input_token_cost"] == 1.25e-05 - assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 - - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_forced_tool_use"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - assert info["prompt_cache_min_tokens"] == 512 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], @@ -280,48 +149,6 @@ def test_fable_5_1_cache_reads_cost_a_quarter_of_fable_5(cost_map): ), model_name -def test_fable_5_1_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - expected_models = { - "global.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1e-05, - "output_cost_per_token": 5e-05, - "cache_creation_input_token_cost": 1.25e-05, - "cache_read_input_token_cost": 2.5e-07, - }, - "us.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 2.75e-07, - }, - "eu.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 2.75e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - -def test_fable_5_1_geo_multiplier_without_fast_mode(): - """Fable 5.1 has no fast mode, so a ``fast`` key here would misprice - ``speed='fast'`` requests.""" - model_data = _load_root_cost_map() - assert model_data["claude-fable-5-1"]["provider_specific_entry"] == {"us": 1.1} - - def test_fable_5_1_present_in_bundled_backup(): backup = GetModelCostMap.load_local_model_cost_map() root = _load_root_cost_map() @@ -334,13 +161,6 @@ def test_fable_5_1_registered_for_bedrock_converse(): assert "anthropic.claude-fable-5-1" in BEDROCK_CONVERSE_MODELS -def test_fable_5_1_provider_resolves_via_model_info(local_model_cost_map): - info = litellm.get_model_info(model="claude-fable-5-1") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "model", [ diff --git a/tests/test_litellm/test_claude_haiku_4_5_config.py b/tests/test_litellm/test_claude_haiku_4_5_config.py index 8755e5d156f..9172b6479a5 100644 --- a/tests/test_litellm/test_claude_haiku_4_5_config.py +++ b/tests/test_litellm/test_claude_haiku_4_5_config.py @@ -7,57 +7,6 @@ import json import os -def test_bedrock_haiku_4_5_configuration(): - """Test that all Bedrock Claude Haiku 4.5 models use bedrock_converse provider""" - # Load model configuration - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - # All Bedrock Haiku 4.5 variants that should use bedrock_converse - bedrock_haiku_models = [ - "anthropic.claude-haiku-4-5-20251001-v1:0", - "anthropic.claude-haiku-4-5@20251001", - "us.anthropic.claude-haiku-4-5-20251001-v1:0", - "eu.anthropic.claude-haiku-4-5-20251001-v1:0", - "apac.anthropic.claude-haiku-4-5-20251001-v1:0", - "jp.anthropic.claude-haiku-4-5-20251001-v1:0", - "global.anthropic.claude-haiku-4-5-20251001-v1:0", - "au.anthropic.claude-haiku-4-5-20251001-v1:0", - ] - - for model in bedrock_haiku_models: - assert model in model_data, f"Model {model} not found in config" - model_info = model_data[model] - - # Verify uses bedrock_converse (not legacy bedrock provider) - assert ( - model_info["litellm_provider"] == "bedrock_converse" - ), f"{model} should use bedrock_converse provider, got {model_info['litellm_provider']}" - - # Verify supports vision (key missing capability) - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - # Verify core capabilities - assert model_info.get("supports_computer_use") is True - assert model_info.get("supports_function_calling") is True - assert model_info.get("supports_tool_choice") is True - assert model_info.get("supports_prompt_caching") is True - assert model_info.get("supports_response_schema") is True - assert model_info.get("supports_pdf_input") is True - assert model_info.get("supports_assistant_prefill") is True - assert model_info.get("supports_reasoning") is True - - # Verify token limits - assert model_info["max_input_tokens"] == 200000 - assert model_info["max_output_tokens"] == 64000 - assert model_info["mode"] == "chat" - - def test_bedrock_haiku_4_5_matches_sonnet_capabilities(): """ Test that Haiku 4.5 has same capabilities as Sonnet 4.5 @@ -97,36 +46,3 @@ def test_bedrock_haiku_4_5_matches_sonnet_capabilities(): assert haiku_info.get(capability) == sonnet_info.get( capability ), f"Capability {capability} mismatch: Haiku={haiku_info.get(capability)}, Sonnet={sonnet_info.get(capability)}" - - -def test_anthropic_api_haiku_4_5_configuration(): - """Test that Anthropic API Claude Haiku 4.5 has correct configuration""" - # Load model configuration - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - # Anthropic API models (not Bedrock) - anthropic_models = [ - "claude-haiku-4-5-20251001", - "claude-haiku-4-5", - ] - - for model in anthropic_models: - assert model in model_data, f"Model {model} not found in config" - model_info = model_data[model] - - # Should use anthropic provider (not bedrock) - assert ( - model_info["litellm_provider"] == "anthropic" - ), f"{model} should use anthropic provider" - - # Should support vision - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - # Should have larger output token limit (64K for Anthropic API) - assert model_info["max_output_tokens"] == 64000 diff --git a/tests/test_litellm/test_claude_opus_4_6_config.py b/tests/test_litellm/test_claude_opus_4_6_config.py index 89d2cd916e0..9a8632924f2 100644 --- a/tests/test_litellm/test_claude_opus_4_6_config.py +++ b/tests/test_litellm/test_claude_opus_4_6_config.py @@ -71,125 +71,6 @@ def test_claude_4_6_australia_region_uses_au_prefix_not_apac(): ), "apac.anthropic.claude-sonnet-4-6 should not be in bedrock_converse_models" -def test_opus_4_6_model_pricing_and_capabilities(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - expected_models = { - "claude-opus-4-6": { - "provider": "anthropic", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "claude-opus-4-6-20260205": { - "provider": "anthropic", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "anthropic.claude-opus-4-6-v1": { - "provider": "bedrock_converse", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "vertex_ai/claude-opus-4-6": { - "provider": "vertex_ai-anthropic_models", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "azure_ai/claude-opus-4-6": { - "provider": "azure_ai", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - } - - for model_name, config in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == config["provider"] - assert info["mode"] == "chat" - assert info["max_input_tokens"] == config["max_input_tokens"] - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_read_input_token_cost"] == 5e-07 - - if config["has_long_context_pricing"]: - assert info["input_cost_per_token_above_200k_tokens"] == 1e-05 - assert info["output_cost_per_token_above_200k_tokens"] == 3.75e-05 - assert info["cache_creation_input_token_cost_above_200k_tokens"] == 1.25e-05 - assert info["cache_read_input_token_cost_above_200k_tokens"] == 1e-06 - else: - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - assert "cache_creation_input_token_cost_above_200k_tokens" not in info - assert "cache_read_input_token_cost_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_opus_4_6_bedrock_regional_model_pricing(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - expected_models = { - "global.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_read_input_token_cost": 5e-07, - }, - "us.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "eu.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "au.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - assert info["supports_assistant_prefill"] is False - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - assert "cache_creation_input_token_cost_above_200k_tokens" not in info - assert "cache_read_input_token_cost_above_200k_tokens" not in info - for key, value in expected.items(): - assert info[key] == value - - def test_opus_4_6_alias_and_dated_metadata_match(): json_path = os.path.join( os.path.dirname(__file__), "../../model_prices_and_context_window.json" diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py index 760512ad31b..e75fdba54ed 100644 --- a/tests/test_litellm/test_claude_opus_4_8_config.py +++ b/tests/test_litellm/test_claude_opus_4_8_config.py @@ -16,7 +16,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -29,102 +28,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_opus_4_8_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = { - "claude-opus-4-8": { - "provider": "anthropic", - "max_input_tokens": 1000000, - }, - "anthropic.claude-opus-4-8": { - "provider": "bedrock_converse", - "max_input_tokens": 1000000, - }, - "vertex_ai/claude-opus-4-8": { - "provider": "vertex_ai-anthropic_models", - "max_input_tokens": 1000000, - }, - "azure_ai/claude-opus-4-8": { - "provider": "azure_ai", - "max_input_tokens": 1000000, - }, - } - - for model_name, config in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == config["provider"] - assert info["mode"] == "chat" - assert info["max_input_tokens"] == config["max_input_tokens"] - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Base pricing matches Opus 4.7: $5 / $25 per MTok, with the standard - # 1.25x cache-write and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_read_input_token_cost"] == 5e-07 - - # Opus 4.x flagships are flat-rate across the full context window. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - assert model_data["claude-opus-4-8"]["supports_native_structured_output"] is True - - -def test_opus_4_8_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - # Global endpoints use base pricing; regional endpoints carry a 10% premium. - expected_models = { - "global.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_read_input_token_cost": 5e-07, - }, - "us.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "eu.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "au.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - def test_opus_4_8_fast_mode_multiplier(): """Opus 4.8 dropped fast-mode pricing to 2x base ($10/$50 per MTok); Opus 4.7 was 6x ($30/$150).""" @@ -134,44 +37,10 @@ def test_opus_4_8_fast_mode_multiplier(): assert entry["fast"] == 2.0 -def test_opus_4_8_present_in_bundled_backup(): - """The bundled backup is the runtime fallback (and what tests load with - ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as - the root cost map, otherwise the model resolves on one path but not the - other.""" - backup = GetModelCostMap.load_local_model_cost_map() - for model_name in ( - "claude-opus-4-8", - "anthropic.claude-opus-4-8", - "global.anthropic.claude-opus-4-8", - "us.anthropic.claude-opus-4-8", - "eu.anthropic.claude-opus-4-8", - "au.anthropic.claude-opus-4-8", - "vertex_ai/claude-opus-4-8", - "vertex_ai/claude-opus-4-8@default", - "azure_ai/claude-opus-4-8", - ): - assert model_name in backup, f"Missing from backup cost map: {model_name}" - assert backup["claude-opus-4-8"]["supports_native_structured_output"] is True - - def test_opus_4_8_registered_for_bedrock_converse(): assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS -def test_opus_4_8_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-opus-4-8`` must resolve to provider ``anthropic``. - - Before the cost-map entry existed, the model was unknown to LiteLLM, so it - could not be tied to the ``anthropic`` provider and an ``anthropic/*`` - wildcard deployment would not match it. - """ - info = litellm.get_model_info(model="claude-opus-4-8") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py index 34744aad17b..285d556ef2b 100644 --- a/tests/test_litellm/test_claude_opus_5_config.py +++ b/tests/test_litellm/test_claude_opus_5_config.py @@ -17,7 +17,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -52,91 +51,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_opus_5_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_providers = { - "claude-opus-5": "anthropic", - "anthropic.claude-opus-5": "bedrock_converse", - "vertex_ai/claude-opus-5": "vertex_ai-anthropic_models", - "azure_ai/claude-opus-5": "azure_ai", - } - - for model_name, provider in expected_providers.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard - # 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05 - assert info["cache_read_input_token_cost"] == 5e-07 - - # Flat rate across the full 1M window, no long-context premium. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no - # assistant prefill. - assert info["supports_adaptive_thinking"] is True - assert info["supports_reasoning"] is True - assert info["supports_sampling_params"] is False - assert info["supports_assistant_prefill"] is False - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_opus_5_bedrock_regional_pricing(): - """Global/base endpoints use base pricing; the us./eu./au./jp. regional - cross-region inference profiles carry a 10% premium.""" - model_data = _load_root_cost_map() - - base_pricing = { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_read_input_token_cost": 5e-07, - } - regional_pricing = { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - } - - expected = { - "anthropic.claude-opus-5": base_pricing, - "global.anthropic.claude-opus-5": base_pricing, - "us.anthropic.claude-opus-5": regional_pricing, - "eu.anthropic.claude-opus-5": regional_pricing, - "au.anthropic.claude-opus-5": regional_pricing, - "jp.anthropic.claude-opus-5": regional_pricing, - } - - for model_name, pricing in expected.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - for key, value in pricing.items(): - assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" - - @pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name): """Bedrock accepts every effort level for Opus 5, so no clamp belongs here. @@ -216,18 +130,6 @@ def test_opus_5_registered_for_bedrock_converse(): assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS -def test_opus_5_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-opus-5`` must resolve to provider ``anthropic``. - - Without the cost-map entry the model is unknown to LiteLLM, so it cannot be - tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment - would not match it.""" - info = litellm.get_model_info(model="claude-opus-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/test_litellm/test_claude_sonnet_5_config.py index 8504326cd21..8c6d2cd1851 100644 --- a/tests/test_litellm/test_claude_sonnet_5_config.py +++ b/tests/test_litellm/test_claude_sonnet_5_config.py @@ -15,7 +15,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -41,96 +40,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_sonnet_5_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_providers = { - "claude-sonnet-5": "anthropic", - "anthropic.claude-sonnet-5": "bedrock_converse", - "vertex_ai/claude-sonnet-5": "vertex_ai-anthropic_models", - "azure_ai/claude-sonnet-5": "azure_ai", - } - - for model_name, provider in expected_providers.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Introductory Sonnet 5 pricing through 2026-08-31: $2 / $10 per MTok, - # with the 1.25x cache-write and 0.1x cache-read multipliers. On - # 2026-09-01 flip these five fields back to the sticker rate, here and - # in both cost-map JSON files (all ten claude-sonnet-5 entries): - # input_cost_per_token: 3e-06 - # output_cost_per_token: 1.5e-05 - # cache_creation_input_token_cost: 3.75e-06 - # cache_creation_input_token_cost_above_1hr: 6e-06 - # cache_read_input_token_cost: 3e-07 - # Regional Bedrock profiles (us./eu./au./jp.) stay at 1.1x those values: - # 3.3e-06 / 1.65e-05 / 4.125e-06 / 6.6e-06 / 3.3e-07 (see - # test_sonnet_5_bedrock_regional_pricing below). - assert info["input_cost_per_token"] == 2e-06 - assert info["output_cost_per_token"] == 1e-05 - assert info["cache_creation_input_token_cost"] == 2.5e-06 - assert info["cache_creation_input_token_cost_above_1hr"] == 4e-06 - assert info["cache_read_input_token_cost"] == 2e-07 - - # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no - # assistant prefill. - assert info["supports_adaptive_thinking"] is True - assert info["supports_reasoning"] is True - assert info["supports_sampling_params"] is False - assert info["supports_assistant_prefill"] is False - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_sonnet_5_bedrock_regional_pricing(): - """Global/base endpoints use base pricing; the us./eu./au./jp. regional - cross-region inference profiles carry a 10% premium.""" - model_data = _load_root_cost_map() - - base_pricing = { - "input_cost_per_token": 2e-06, - "output_cost_per_token": 1e-05, - "cache_creation_input_token_cost": 2.5e-06, - "cache_creation_input_token_cost_above_1hr": 4e-06, - "cache_read_input_token_cost": 2e-07, - } - regional_pricing = { - "input_cost_per_token": 2.2e-06, - "output_cost_per_token": 1.1e-05, - "cache_creation_input_token_cost": 2.75e-06, - "cache_creation_input_token_cost_above_1hr": 4.4e-06, - "cache_read_input_token_cost": 2.2e-07, - } - - expected = { - "anthropic.claude-sonnet-5": base_pricing, - "global.anthropic.claude-sonnet-5": base_pricing, - "us.anthropic.claude-sonnet-5": regional_pricing, - "eu.anthropic.claude-sonnet-5": regional_pricing, - "au.anthropic.claude-sonnet-5": regional_pricing, - "jp.anthropic.claude-sonnet-5": regional_pricing, - } - - for model_name, pricing in expected.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in pricing.items(): - assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" - - def test_sonnet_5_present_in_bundled_backup(): """The bundled backup is the runtime fallback (and what tests load with ``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the @@ -144,18 +53,6 @@ def test_sonnet_5_registered_for_bedrock_converse(): assert "anthropic.claude-sonnet-5" in BEDROCK_CONVERSE_MODELS -def test_sonnet_5_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-sonnet-5`` must resolve to provider ``anthropic``. - - Before the cost-map entry existed, the model was unknown to LiteLLM, so it - could not be tied to the ``anthropic`` provider and an ``anthropic/*`` - wildcard deployment would not match it.""" - info = litellm.get_model_info(model="claude-sonnet-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py index e33bcfb8378..4e770be7c3e 100644 --- a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py +++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py @@ -27,17 +27,6 @@ BACKUP_MAP = os.path.join( ) -@pytest.fixture(autouse=True) -def _use_local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - try: - yield - finally: - litellm.model_cost = original_model_cost - - def _load(path: str) -> dict: with open(path, encoding="utf-8") as f: return json.load(f) @@ -47,50 +36,6 @@ def _cloudflare_keys(data: dict) -> set: return {k for k in data if k.startswith("cloudflare/")} -def test_glm_5_2_entry_is_present_and_well_formed(): - entry = litellm.model_cost["cloudflare/@cf/zai-org/glm-5.2"] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["input_cost_per_token"] > 0 - assert entry["output_cost_per_token"] > 0 - - -def test_vision_model_is_flagged_supports_vision(): - entry = litellm.model_cost["cloudflare/@cf/meta/llama-3.2-11b-vision-instruct"] - assert entry["litellm_provider"] == "cloudflare" - assert entry.get("supports_vision") is True - - -def test_additional_current_models_are_present(): - for key in ( - "cloudflare/@cf/openai/gpt-oss-120b", - "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast", - ): - entry = litellm.model_cost[key] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["input_cost_per_token"] > 0 - assert entry["output_cost_per_token"] > 0 - - -@pytest.mark.parametrize( - "key, published_price_per_audio_minute", - [ - ("cloudflare/@cf/openai/whisper", 0.00045), - ("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051), - ], -) -def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute): - entry = litellm.model_cost[key] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "audio_transcription" - assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] - assert entry["output_cost_per_second"] == 0.0 - assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60) - - def test_root_and_backup_have_identical_cloudflare_keys(): if not os.path.exists(ROOT_MAP): pytest.skip("root cost map only ships in source checkouts") diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/test_litellm/test_daybreak_model_metadata.py index dbb7ecdffac..c3bac14dbbd 100644 --- a/tests/test_litellm/test_daybreak_model_metadata.py +++ b/tests/test_litellm/test_daybreak_model_metadata.py @@ -32,22 +32,6 @@ def _load(path): return json.load(f) -@pytest.mark.parametrize("model", DAYBREAK_MODELS) -def test_daybreak_capability_contract(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "openai" - assert info["mode"] == "chat" - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] - - assert info["supports_computer_use"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - - def test_blue_alias_matches_its_snapshot_computer_use(): cost_map = _load(MAIN_PATH) diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/test_litellm/test_deepseek_model_metadata.py index b9eb33f0972..9cbd14ebd1e 100644 --- a/tests/test_litellm/test_deepseek_model_metadata.py +++ b/tests/test_litellm/test_deepseek_model_metadata.py @@ -39,26 +39,6 @@ class TestDeepSeekModelCostEntries: """Verify that provider-prefixed DeepSeek entries contain the same capability flags as their bare-name counterparts in the JSON files.""" - def test_deepseek_chat_supports_response_schema_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_response_schema") is True - - def test_deepseek_reasoner_supports_response_schema_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_response_schema") is True - - def test_deepseek_chat_supports_system_messages_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_system_messages") is True - - def test_deepseek_reasoner_supports_system_messages_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_system_messages") is True - def test_deepseek_chat_max_input_tokens_matches_bare_in_backup(self): data = _load_backup_json() bare = data.get("deepseek-chat", {}) @@ -71,26 +51,6 @@ class TestDeepSeekModelCostEntries: prefixed = data.get("deepseek/deepseek-reasoner", {}) assert prefixed.get("max_output_tokens") == bare.get("max_output_tokens") - def test_main_json_deepseek_chat_supports_response_schema(self): - main_path = os.path.join( - os.path.dirname(os.path.dirname(litellm.__file__)), - "model_prices_and_context_window.json", - ) - with open(main_path, encoding="utf-8") as f: - data = json.load(f) - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_response_schema") is True - - def test_main_json_deepseek_reasoner_supports_response_schema(self): - main_path = os.path.join( - os.path.dirname(os.path.dirname(litellm.__file__)), - "model_prices_and_context_window.json", - ) - with open(main_path, encoding="utf-8") as f: - data = json.load(f) - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_response_schema") is True - # --------------------------------------------------------------------------- # API-level tests – verify supports_response_schema returns True diff --git a/tests/test_litellm/test_exception_header_preservation.py b/tests/test_litellm/test_exception_header_preservation.py index 6ea478c633b..dd142d9d40b 100644 --- a/tests/test_litellm/test_exception_header_preservation.py +++ b/tests/test_litellm/test_exception_header_preservation.py @@ -18,6 +18,7 @@ from litellm.exceptions import ( ImageFetchError, MidStreamFallbackError, RateLimitError, + ServiceUnavailableError, ) @@ -312,3 +313,85 @@ class TestProxyHeaderExtraction: # Verify headers are extracted and prefixed correctly assert headers.get("llm_provider-x-request-id") == "req-abc123" assert headers.get("llm_provider-x-ms-region") == "eastus" + + +class TestBedrockErrorHeaders: + """A BedrockError built with headers but no response still exposes them (LIT-5428).""" + + def test_synthesized_response_carries_headers(self): + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="Amazon Bedrock is unable to process your request.", + headers={"x-amzn-RequestId": "req-base-500"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-base-500" + assert str(error.request.url) == str(BedrockError(status_code=500, message="boom").request.url) + assert str(error.response.request.url) == str(error.request.url) + + def test_synthesized_response_without_headers_stays_empty(self): + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError(status_code=500, message="boom") + + assert dict(error.response.headers) == {} + + def test_explicit_response_is_kept(self): + from litellm.llms.bedrock.common_utils import BedrockError + + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "from-response"}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + error = BedrockError( + status_code=500, + message="boom", + headers={"x-amzn-RequestId": "from-headers"}, + response=provider_response, + ) + + assert error.response is provider_response + + def test_proxy_extraction_surfaces_bedrock_request_id(self): + """End-to-end shape the proxy error handler returns to the caller.""" + from litellm.litellm_core_utils.exception_mapping_utils import exception_type + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers, + ) + from litellm.llms.bedrock.common_utils import BedrockError + + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-proxy-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + with pytest.raises(ServiceUnavailableError) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=BedrockError( + status_code=500, + message=provider_response.text, + headers=provider_response.headers, + response=provider_response, + ), + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + # Mirrors ProxyBaseLLMRequestProcessing._handle_llm_api_exception + error = exc_info.value + headers = getattr(error, "headers", None) or {} + if not headers: + _response = getattr(error, "response", None) + if _response is not None: + _response_headers = getattr(_response, "headers", None) + if _response_headers: + headers = get_response_headers(dict(_response_headers)) + + assert headers.get("llm_provider-x-amzn-requestid") == "req-proxy-500" diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index a7a9e0fc37d..701938f5677 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -14,26 +14,9 @@ import os import pytest -import litellm from litellm.utils import get_model_info -@pytest.fixture(scope="module", autouse=True) -def _local_model_cost_map(): - """ - Point litellm at the bundled cost map for the duration of this module - only. ``mp.undo()`` restores both the environment variable and - ``litellm.model_cost`` so nothing leaks into later tests. - """ - mp = pytest.MonkeyPatch() - mp.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - mp.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - get_model_info.cache_clear() - yield - mp.undo() - get_model_info.cache_clear() - - NEW_ENTRIES = { "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "input_cost_per_token": 1.32e-06, @@ -54,19 +37,6 @@ def model_data(): return json.load(f) -def test_fireworks_serverless_entries_exist(model_data): - """The new prefixed entry carries the pricing and metadata from #37274.""" - for key, expected in NEW_ENTRIES.items(): - assert key in model_data, f"{key} is missing from model_prices_and_context_window.json" - entry = model_data[key] - for field, value in expected.items(): - assert entry[field] == pytest.approx(value), f"{key}.{field}" - assert entry["litellm_provider"] == "fireworks_ai" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["supports_vision"] is False - - def test_bare_fireworks_ids_resolve_through_prefixed_entries(): """Bare IDs from #37274 resolve via the provider-prefix lookup path.""" for bare_id, prefixed_key in [ diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py index a60fa9466e6..e07efbcc913 100644 --- a/tests/test_litellm/test_gpt_5_5_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py @@ -1,54 +1,6 @@ import json from pathlib import Path -import pytest - -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - -@pytest.mark.parametrize("model", ["azure_ai/gpt-5.5", "azure_ai/gpt-5.5-2026-04-23"]) -def test_azure_ai_gpt_5_5_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "azure_ai" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 3e-05 - assert info["cache_read_input_token_cost"] == 5e-07 - - assert info["input_cost_per_token_above_272k_tokens"] == 1e-05 - assert info["output_cost_per_token_above_272k_tokens"] == 4.5e-05 - assert info["cache_read_input_token_cost_above_272k_tokens"] == 1e-06 - - assert info["input_cost_per_token_priority"] == 1e-05 - assert info["output_cost_per_token_priority"] == 6e-05 - - assert info["max_input_tokens"] == 1050000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_web_search"] is True - # gpt-5.5 dropped minimal reasoning effort support (true on gpt-5.4) - assert info["supports_minimal_reasoning_effort"] is False - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "azure_ai" - def test_azure_ai_gpt_5_5_backup_matches_main(): """Ensure the bundled model cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py index 314fd63c4cc..8c41e474486 100644 --- a/tests/test_litellm/test_gpt_realtime_mode.py +++ b/tests/test_litellm/test_gpt_realtime_mode.py @@ -1,10 +1,8 @@ import json from pathlib import Path -import pytest from typing_extensions import get_args, get_type_hints -import litellm from litellm.types.utils import ModelInfoBase REALTIME_ONLY_GPT_MODELS = ( @@ -43,44 +41,11 @@ REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS = ( ALL_REALTIME_ONLY_GPT_MODELS = REALTIME_ONLY_GPT_MODELS + REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS -def _load_cost_map() -> dict: - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - return json.load(f) - - def test_realtime_is_a_valid_mode_literal(): hints = get_type_hints(ModelInfoBase, include_extras=False) assert "realtime" in get_args(hints["mode"]) -@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS) -def test_realtime_only_gpt_models_are_mode_realtime(model): - """These models only serve /v1/realtime and are rejected by /v1/chat/completions - ("This is not a chat model ..."), so they must not be tagged mode=chat.""" - info = _load_cost_map()[model] - assert info["supported_endpoints"] == ["/v1/realtime"] - assert info["mode"] == "realtime" - - -@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS) -def test_realtime_only_gpt_4o_models_are_mode_realtime(model): - """gpt-4o(-mini)-realtime-preview are realtime-only and must not be mode=chat.""" - assert _load_cost_map()[model]["mode"] == "realtime" - - -def test_get_model_info_reports_realtime_mode(monkeypatch): - """get_model_info must resolve the retag against the bundled cost map, not the - hosted map fetched from main, which lags this repo until the next promotion.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() - try: - assert litellm.get_model_info("gpt-realtime-mini")["mode"] == "realtime" - finally: - litellm.get_model_info.cache_clear() - - def test_backup_matches_main_for_realtime_models(): repo_root = Path(__file__).parents[2] with open(repo_root / "model_prices_and_context_window.json") as f: diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py index 6f1ba702d8d..d73311baae9 100644 --- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py +++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py @@ -3,8 +3,6 @@ from pathlib import Path import pytest -import litellm -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider REPO_ROOT = Path(__file__).parents[2] MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -27,56 +25,6 @@ def _load(path): return json.load(f) - -@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS) -def test_medium_3_5_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - assert info["max_tokens"] == 262144 - - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_assistant_prefill"] is True - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "mistral" - - -def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map): - """LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must - return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values.""" - info = litellm.get_model_info(model="mistral/mistral-medium-latest") - - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - assert info["max_input_tokens"] == 262144 - assert info["supports_reasoning"] is True - - -def test_mistral_medium_2508_keeps_medium_3_1_specs(): - """The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing.""" - info = _load(MAIN_PATH).get("mistral/mistral-medium-2508") - assert info is not None, "mistral/mistral-medium-2508 missing from cost map" - - assert info["input_cost_per_token"] == 4e-07 - assert info["output_cost_per_token"] == 2e-06 - assert info["max_input_tokens"] == 131072 - assert info.get("supports_reasoning") is not True - - @pytest.mark.parametrize("model", SYNCED_MODELS) def test_backup_matches_main(model): """Ensure the bundled (backup) cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py index 0442321ba0b..182c444bac9 100644 --- a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py +++ b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py @@ -18,29 +18,6 @@ def _load(path): return json.load(f) -@pytest.mark.parametrize("model", SMALL_4_0_MODELS) -def test_small_4_0_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.5e-07 - assert info["output_cost_per_token"] == 6e-07 - - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - assert info["max_tokens"] == 262144 - - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_assistant_prefill"] is True - - @pytest.mark.parametrize("model", SMALL_4_0_MODELS) def test_backup_matches_main(model): main_cost = _load(MAIN_PATH) diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py index 0587883aa44..02527a98711 100644 --- a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py @@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) - -@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) -def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): - info = _load_cost_map().get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "meta" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cached_cost - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_pdf_input"] is True - assert info["supports_web_search"] is True - assert info["supports_minimal_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] - assert info["supported_modalities"] == ["text", "image", "video"] - assert info["supported_output_modalities"] == ["text"] - - assert info["search_context_cost_per_query"] == { - "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, - } - - @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_2_cost_per_token( local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py index 1ecd9490f78..92b099fc780 100644 --- a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py @@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) - -@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) -def test_muse_spark_1_3_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): - info = _load_cost_map().get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "meta" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cached_cost - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_pdf_input"] is True - assert info["supports_web_search"] is True - assert info["supports_minimal_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] - assert info["supported_modalities"] == ["text", "image", "video"] - assert info["supported_output_modalities"] == ["text"] - - assert info["search_context_cost_per_query"] == { - "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, - } - - @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_3_cost_per_token( local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float diff --git a/tests/test_litellm/test_replicate_model_key_format.py b/tests/test_litellm/test_replicate_model_key_format.py index 8c2b72f8ed2..77ae5e1b069 100644 --- a/tests/test_litellm/test_replicate_model_key_format.py +++ b/tests/test_litellm/test_replicate_model_key_format.py @@ -20,14 +20,6 @@ def test_replicate_models_have_valid_key_prefix(model_cost: dict[str, Any]) -> N ) -def test_replicate_openai_gpt_oss_20b_key_exists(model_cost: dict[str, Any]) -> None: - assert "replicate/openai/gpt-oss-20b" in model_cost - info = model_cost["replicate/openai/gpt-oss-20b"] - assert info["litellm_provider"] == "replicate" - assert info["mode"] == "chat" - assert info["supports_function_calling"] is True - - def test_replicate_backup_matches_main() -> None: repo_root = Path(__file__).parents[2] main_path = repo_root / "model_prices_and_context_window.json" diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py index c9e2863d240..b9764eca2f8 100644 --- a/tests/test_litellm/test_together_ai_model_metadata.py +++ b/tests/test_litellm/test_together_ai_model_metadata.py @@ -5,7 +5,6 @@ from typing import Final import pytest from pydantic import TypeAdapter -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider REPO_ROOT: Final = Path(__file__).parents[2] @@ -77,59 +76,6 @@ def cost_map() -> CostMap: return COST_MAP_ADAPTER.validate_python(json.load(f)) -@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS) -def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str): - info = cost_map.get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - assert info["litellm_provider"] == "together_ai" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] >= 0 - assert info["output_cost_per_token"] >= info["input_cost_per_token"] - assert "deprecation_date" not in info - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.removeprefix("together_ai/") - assert provider == "together_ai" - - -def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap): - info = cost_map["together_ai/moonshotai/Kimi-K3"] - assert info["input_cost_per_token"] == 3e-06 - assert info["output_cost_per_token"] == 1.5e-05 - assert info["max_input_tokens"] == 1048576 - assert info["supports_function_calling"] is True - assert info["supports_tool_choice"] is True - assert info["supports_response_schema"] is True - assert info["supports_vision"] is True - assert info["supports_reasoning"] is True - - -def test_together_glm_52_pricing(cost_map: CostMap): - info = cost_map["together_ai/zai-org/GLM-5.2"] - assert info["input_cost_per_token"] == 1.4e-06 - assert info["output_cost_per_token"] == 4.4e-06 - assert info["max_input_tokens"] == 1048575 - assert info["max_output_tokens"] == 128000 - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - - -def test_together_glm_53_flash_pricing_and_capabilities(cost_map: CostMap): - info = cost_map["together_ai/zai-org/GLM-5.3-Flash"] - assert info["input_cost_per_token"] == 1.5e-07 - assert info["output_cost_per_token"] == 5e-07 - assert info["cache_read_input_token_cost"] == 3e-08 - assert info["max_input_tokens"] == 1048575 - assert info["max_output_tokens"] == 128000 - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_response_schema"] is True - assert info["supports_vision"] is True - assert info["supports_reasoning"] is True - - def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost_map: CostMap): inflated = sorted( model @@ -142,21 +88,6 @@ def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost assert inflated == [] -def test_together_multilingual_e5_embedding_entry(cost_map: CostMap): - info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] - assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == 2e-08 - assert info["max_input_tokens"] == 514 - assert info["output_vector_size"] == 1024 - - -def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap): - info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"] - assert info["input_cost_per_token"] == 1.04e-06 - assert info["output_cost_per_token"] == 1.04e-06 - assert info["max_input_tokens"] == 131072 - - @pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS)) def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str): info = cost_map.get(model) @@ -210,32 +141,7 @@ CACHED_INPUT_MODELS: Final = ( ) -@pytest.mark.parametrize("model", CACHED_INPUT_MODELS) -def test_together_cached_input_model_carries_cache_read_pricing(cost_map: CostMap, model: str): - info = cost_map.get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - assert info.get("supports_prompt_caching") is True - cache_read = info.get("cache_read_input_token_cost") - assert isinstance(cache_read, float) - assert 0 < cache_read < info["input_cost_per_token"] - assert "cache_creation_input_token_cost" not in info - - def test_together_prompt_caching_flag_implies_cache_read_rate(cost_map: CostMap): for model, info in cost_map.items(): if model.startswith("together_ai/") and info.get("supports_prompt_caching"): assert "cache_read_input_token_cost" in info, f"{model} flags caching without a cache read rate" - - -def test_together_deepseek_v4_flash_cache_read_rate(cost_map: CostMap): - info = cost_map["together_ai/deepseek-ai/DeepSeek-V4-Flash-0731"] - assert info["input_cost_per_token"] == 1.4e-07 - assert info["cache_read_input_token_cost"] == 3e-08 - assert info["output_cost_per_token"] == 2.8e-07 - - -def test_together_qwen_37_max_repriced_to_current_together_rate(cost_map: CostMap): - info = cost_map["together_ai/Qwen/Qwen3.7-Max"] - assert info["input_cost_per_token"] == 2.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - assert info["cache_read_input_token_cost"] == 5e-07 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8a56a84ade7..4ce68b71079 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -6092,3 +6092,47 @@ class TestFinalOptionalParamsLineRedaction: assert "'max_tokens': 17" in printed assert "'temperature': 0.25" in printed + + +def _credential_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [record.getMessage() for record in caplog.records if "litellm_credential_name=" in record.getMessage()] + + +def test_load_credentials_from_list_warns_when_the_named_credential_is_not_loaded( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.utils import load_credentials_from_list + + monkeypatch.setattr(litellm, "credential_list", []) + request_kwargs = {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + assert _credential_warnings(caplog) == [ + "litellm_credential_name=openai-cred matched none of the 0 loaded credentials; the request runs without it" + ] + + +def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_without_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.types.utils import CredentialItem + from litellm.utils import load_credentials_from_list + + loaded = CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db", "api_base": "https://credential.example"}, + credential_info={}, + ) + monkeypatch.setattr(litellm, "credential_list", [loaded]) + request_kwargs = {"litellm_credential_name": "openai-cred", "api_base": "https://request.example"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == { + "litellm_credential_name": "openai-cred", + "api_base": "https://request.example", + "api_key": "sk-from-db", + } + assert _credential_warnings(caplog) == [] diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py new file mode 100644 index 00000000000..02d274bb405 --- /dev/null +++ b/tests/test_litellm_rust/conftest.py @@ -0,0 +1,24 @@ +import os + +import pytest + + +def pytest_collection_modifyitems(items): + rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + if not rust_enabled: + skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") + for item in items: + item.add_marker(skip) + return + + try: + from litellm.rust_bridge import _native # noqa: F401 # validates the installed extension + except ImportError as error: + raise pytest.UsageError( + "LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension" + ) from error diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py new file mode 100644 index 00000000000..d5b1fce1139 --- /dev/null +++ b/tests/test_litellm_rust/test_ocr.py @@ -0,0 +1,72 @@ +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +import litellm + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + requests.append( + { + "headers": {name.lower(): value for name, value in self.headers.items()}, + "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))), + } + ) + if self.headers.get("User-Agent", "").startswith("python-httpx"): + self.send_response(418) + self.end_headers() + return + response = json.dumps( + { + "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, + } + ).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + + def log_message(self, format, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + try: + yield server, requests + finally: + server.shutdown() + server.server_close() + thread.join() + + +def test_ocr_with_rust_extension(ocr_server): + server, requests = ocr_server + host, port = server.server_address + + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + api_key="test-key", + api_base=f"http://{host}:{port}", + ) + + assert response.pages[0].markdown == "native OCR response" + assert len(requests) == 1 + assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") + assert requests[0]["body"] == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + } diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..0c0952289e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22180 + "limit": 22174 }, "LIT002": { - "limit": 26729 + "limit": 26715 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16426 + "limit": 16398 }, "LIT011": { - "limit": 5506 + "limit": 5504 }, "LIT012": { "limit": 4486 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts index 442dcd48f66..e0a5271366f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts @@ -313,6 +313,26 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { return agentData; }; +export const parseMcpPermissionsForForm = (agent: any) => ({ + allowed_mcp_servers_and_groups: { + servers: agent.object_permission?.mcp_servers ?? [], + accessGroups: agent.object_permission?.mcp_access_groups ?? [], + toolsets: agent.object_permission?.mcp_toolsets ?? [], + }, + mcp_tool_permissions: agent.object_permission?.mcp_tool_permissions ?? {}, +}); + +/** + * Always includes every MCP key (empty when cleared) so removals persist; + * the proxy merges object_permission per key, leaving non-MCP grants untouched. + */ +export const buildMcpObjectPermission = (values: any) => ({ + mcp_servers: values.allowed_mcp_servers_and_groups?.servers ?? [], + mcp_access_groups: values.allowed_mcp_servers_and_groups?.accessGroups ?? [], + mcp_toolsets: values.allowed_mcp_servers_and_groups?.toolsets ?? [], + mcp_tool_permissions: values.mcp_tool_permissions ?? {}, +}); + /** * Parse agent data for form fields */ @@ -356,5 +376,6 @@ export const parseAgentForForm = (agent: any) => { : [], // extra_headers: already an array of strings extra_headers: agent.extra_headers ?? [], + ...parseMcpPermissionsForForm(agent), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index de1eb153f6f..357f924cbe7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -1,4 +1,5 @@ import React from "react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; @@ -10,6 +11,12 @@ vi.mock("@/components/networking", () => ({ getAgentInfo: vi.fn(), patchAgentCall: vi.fn(), getAgentCreateMetadata: vi.fn(), + getProxyBaseUrl: vi.fn(() => ""), + getUiConfig: vi.fn(async () => ({})), + fetchMCPServers: vi.fn(async () => []), + fetchMCPAccessGroups: vi.fn(async () => []), + fetchMCPToolsets: vi.fn(async () => []), + listMCPTools: vi.fn(async () => ({ tools: [] })), })); vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ @@ -111,7 +118,14 @@ const bedrockAgentcoreInfo: AgentCreateInfo = { const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); -const renderView = () => render(); +const renderView = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); +}; const openEditor = async (user: ReturnType) => { await user.click(await screen.findByRole("tab", { name: "Settings" })); @@ -161,6 +175,7 @@ describe("AgentInfoView update payload", () => { rpm_limit: 222, session_tpm_limit: 333, session_rpm_limit: 444, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); @@ -201,6 +216,7 @@ describe("AgentInfoView update payload", () => { rpm_limit: 222, session_tpm_limit: 333, session_rpm_limit: 444, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); @@ -278,9 +294,30 @@ describe("AgentInfoView update payload", () => { api_base: "https://other.example.com", model: "langgraph/asst_1", }, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); + it("keeps the agent's existing MCP grants in the update payload", async () => { + const existingMcpGrants = { + mcp_servers: ["srv-1"], + mcp_access_groups: ["grp-a"], + mcp_toolsets: ["toolset-1"], + mcp_tool_permissions: { "srv-1": ["tool_x"] }, + }; + vi.mocked(networking.getAgentInfo).mockResolvedValue({ + ...A2A_AGENT, + object_permission: existingMcpGrants, + } as never); + const user = setup(); + renderView(); + await openEditor(user); + + await save(user); + + expect(patchedPayload().object_permission).toEqual(existingMcpGrants); + }); + it("preserves the full AgentCore runtime ARN (including the resource id after runtime/) across an unedited save", async () => { vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]); vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx index 0936b8e13db..29d97a20afe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx @@ -24,6 +24,18 @@ vi.mock("./agent_form_fields", () => ({ unmountedA2AFieldNames: () => [], })); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }), +})); + +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + default: () =>
, +})); + +vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ + default: () =>
, +})); + const agent = { agent_id: "agent-1", agent_name: "support-agent", @@ -62,5 +74,18 @@ describe("AgentInfoView settings", () => { expect(token).toBe("sk-test"); expect(agentId).toBe("agent-1"); expect(payload.tpm_limit).toBe(42); + const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }; + expect(payload.object_permission).toEqual(clearedMcpGrants); + }); + + it("shows MCP grants with server names on the overview tab", async () => { + vi.mocked(networking.getAgentInfo).mockResolvedValue({ + ...agent, + object_permission: { mcp_servers: ["srv-1"] }, + } as unknown as Agent); + + render(); + + expect(await screen.findByText("github (srv-1)")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index eddeeec674b..6cb99e9692f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -15,16 +15,27 @@ import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } import { Agent } from "@/components/agents/types"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import KeyInfoView from "@/components/templates/key_info_view"; +import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; +import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; import AgentVirtualKeys from "./agent_virtual_keys"; import AgentFormFields, { unmountedA2AFieldNames } from "./agent_form_fields"; import DynamicAgentFormFields, { buildDynamicAgentData, unmountedDynamicFieldNames } from "./dynamic_agent_form_fields"; -import { AGENT_FORM_CONFIG, buildAgentDataFromForm, parseAgentForForm } from "./agent_config"; +import { + AGENT_FORM_CONFIG, + buildAgentDataFromForm, + buildMcpObjectPermission, + parseAgentForForm, + parseMcpPermissionsForForm, +} from "./agent_config"; import { AgentFormField, AgentFormValues, AgentNumberInput, AgentRequestPayload, + McpServerSelection, + labelWithHint, omitFieldValues, useCollapsiblePanels, } from "./AgentFormKit"; @@ -111,7 +122,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT } else { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset(parseDynamicAgentForForm(data, typeInfo)); + form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) }); } else { form.reset(parseAgentForForm(data)); } @@ -131,7 +142,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT if (agentType !== "a2a") { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset(parseDynamicAgentForForm(agent, typeInfo)); + form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) }); } } } @@ -139,6 +150,14 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT const selectedAgentTypeInfo = agentTypeMetadata.find((t) => t.agent_type === detectedAgentType); const watchedFormValues = useWatch({ control: form.control }); + const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" }); + const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" }); + const { data: mcpServers = [] } = useMCPServers(); + + const mcpServerLabel = (serverId: string) => { + const server = mcpServers.find((s) => s.server_id === serverId); + return server?.server_name ? `${server.server_name} (${serverId})` : serverId; + }; const discoveryRequest = useMemo( () => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo), @@ -199,7 +218,10 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT ? overlayDiscoveredCardParams(built, appliedDiscoveredSelection.selected_card) : built; - await patchAgentCall(accessToken, agentId, updateData); + await patchAgentCall(accessToken, agentId, { + ...updateData, + object_permission: buildMcpObjectPermission(values), + }); toast.success("Agent updated successfully"); setIsEditing(false); fetchAgentInfo(); @@ -337,13 +359,20 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.object_permission && (agent.object_permission.mcp_servers?.length || agent.object_permission.mcp_access_groups?.length || + agent.object_permission.mcp_toolsets?.length || (agent.object_permission.mcp_tool_permissions && Object.keys(agent.object_permission.mcp_tool_permissions).length > 0)) && (

MCP Tool Permissions

{agent.object_permission.mcp_servers && agent.object_permission.mcp_servers.length > 0 && ( - {agent.object_permission.mcp_servers.join(", ")} + +
+ {agent.object_permission.mcp_servers.map((serverId) => ( +
{mcpServerLabel(serverId)}
+ ))} +
+ )} {agent.object_permission.mcp_access_groups && agent.object_permission.mcp_access_groups.length > 0 && ( @@ -351,13 +380,16 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.object_permission.mcp_access_groups.join(", ")} )} + {agent.object_permission.mcp_toolsets && agent.object_permission.mcp_toolsets.length > 0 && ( + {agent.object_permission.mcp_toolsets.join(", ")} + )} {agent.object_permission.mcp_tool_permissions && Object.keys(agent.object_permission.mcp_tool_permissions).length > 0 && (
{Object.entries(agent.object_permission.mcp_tool_permissions).map(([serverId, tools]) => (
- {serverId}:{" "} + {mcpServerLabel(serverId)}:{" "} {Array.isArray(tools) ? tools.join(", ") : String(tools)}
))} @@ -457,6 +489,41 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {rateLimitField("session_rpm_limit", "Session RPM Limit")}
+ +

MCP Servers

+ + + {({ value, onChange }) => ( + + )} + + +
+ ) => + form.setValue("mcp_tool_permissions", toolPerms) + } + /> +
+
{stats.classifier_cost == null && ( @@ -277,9 +290,10 @@ const BenchmarksBody: React.FC = ({ isPending, error, data,

Compares your actual routed spend with the estimated cost of using only the most expensive model configured in the auto-router. It accounts for both the cache savings from staying on one model and the added cache costs from - switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. The - range counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets - savings by UTC day. + switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. + Classification cost per 1K turns is averaged over all auto-router turns, including those that skip + classification. The range counts whole sessions that overlap it, so totals can differ slightly from the Overall + tab, which buckets savings by UTC day.

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts index 9a2d7a0b0ec..5d2c48e6440 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts @@ -7,6 +7,7 @@ import { SAVINGS_DRIVERS, SAVINGS_SERIES, buildDailyToolSeries, + classificationRatePer1kTurns, computeCacheLeakage, formatRangeLabel, isAnthropicModel, @@ -401,6 +402,23 @@ describe("usd", () => { }); }); +describe("classificationRatePer1kTurns", () => { + it("normalizes total classification cost to one thousand turns", () => { + expect(classificationRatePer1kTurns(342.18, 140815)).toBe("($2.43 / 1K turns)"); + expect(classificationRatePer1kTurns(0.0004, 100)).toBe("($0.0040 / 1K turns)"); + }); + + it("shows a floor instead of rounding a real cost down to zero", () => { + expect(classificationRatePer1kTurns(0.00001, 1000)).toBe("(<$0.0001 / 1K turns)"); + expect(classificationRatePer1kTurns(0.0001, 1000)).toBe("($0.0001 / 1K turns)"); + }); + + it("reports zero when there are no turns or no classification cost", () => { + expect(classificationRatePer1kTurns(0, 0)).toBe("($0.00 / 1K turns)"); + expect(classificationRatePer1kTurns(0, 100)).toBe("($0.00 / 1K turns)"); + }); +}); + describe("savings driver colours", () => { it("keeps a driver's colour when a driver above it is filtered out", () => { // Charts colour by position in the data they are given, and the donut is given diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts index 7019b0d3301..464c779aa2b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts @@ -10,6 +10,13 @@ export const usd = (value: number): string => { return `${value < 0 ? "-" : ""}$${formatNumberWithCommas(magnitude, decimals)}`; }; +export const classificationRatePer1kTurns = (classifierCost: number, turns: number): string => { + if (turns <= 0) return `(${usd(0)} / 1K turns)`; + const rate = (classifierCost * 1000) / turns; + if (rate > 0 && rate < 0.0001) return "(<$0.0001 / 1K turns)"; + return `(${usd(rate)} / 1K turns)`; +}; + export const pct = (ratio: number): string => `${formatNumberWithCommas(ratio * 100, 1)}%`; export const shortDate = (iso: string): string => diff --git a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx index 30ce62d8081..a663e16c2b2 100644 --- a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx @@ -1,43 +1,19 @@ "use client"; -import { Suspense, useEffect } from "react"; -import { useRouter, useSearchParams } from "next/navigation"; +import { Suspense } from "react"; import { useChatShell } from "@/contexts/ChatShellContext"; -import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; -import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface"; // useSearchParams() requires a Suspense boundary for static export. function IntegrationsPageContent() { const { accessToken, selectedMCPServers, setSelectedMCPServers } = useChatShell(); - const router = useRouter(); - const searchParams = useSearchParams(); - const oauthReturn = searchParams.get("mcpOauthReturn"); - // Set by the gateway DCR authorize when a DCR client sends the user here to - // authorize servers before finishing sign-in (see gateway_dcr_flow.py). The - // handle keys the sealed per-flow cookie; connect_client is the client origin - // for display only. connect_flow is NOT cleaned from the URL: the finish form - // needs it, and the sealed cookie (not the URL) is the security boundary. - const connectFlow = searchParams.get("connect_flow"); - const connectClient = searchParams.get("connect_client"); - - // Clean up the OAuth return param after it's been consumed — real routing means - // we no longer need it to pick a tab, but it should not linger in the address bar. - useEffect(() => { - if (oauthReturn) { - const url = new URL(window.location.href); - url.searchParams.delete("mcpOauthReturn"); - router.replace(url.pathname + url.search); - } - }, [oauthReturn, router]); return (
- {connectFlow && } -
); diff --git a/ui/litellm-dashboard/src/app/connect/page.test.tsx b/ui/litellm-dashboard/src/app/connect/page.test.tsx index 7d49a8b6a4c..6a6ba24bd87 100644 --- a/ui/litellm-dashboard/src/app/connect/page.test.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.test.tsx @@ -2,101 +2,29 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { render, screen } from "@testing-library/react"; import ConnectPage from "./page"; -interface PanelProps { +interface SurfaceProps { accessToken: string; selectedServers: string[]; onChange: (servers: string[]) => void; - connectMode?: boolean; } -interface BannerProps { - flowHandle: string; - clientOrigin: string | null; -} - -const { mockReplace, mockPanel, mockBanner, state } = vi.hoisted(() => { - const state = { - oauthReturn: null as string | null, - connectFlow: null as string | null, - connectClient: null as string | null, - }; - return { - state, - mockReplace: vi.fn(), - mockPanel: vi.fn((_props: PanelProps) =>
), - mockBanner: vi.fn((_props: BannerProps) =>
), - }; -}); - -vi.mock("next/navigation", () => ({ - useRouter: () => ({ replace: mockReplace }), - useSearchParams: () => ({ - get: (key: string) => { - if (key === "mcpOauthReturn") return state.oauthReturn; - if (key === "connect_flow") return state.connectFlow; - if (key === "connect_client") return state.connectClient; - return null; - }, - }), +const { mockSurface } = vi.hoisted(() => ({ + mockSurface: vi.fn((_props: SurfaceProps) =>
), })); + vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "token-123" }), })); -vi.mock("@/components/chat/MCPAppsPanel", () => ({ default: mockPanel })); -vi.mock("@/components/chat/ConnectFlowBanner", () => ({ default: mockBanner })); +vi.mock("@/components/chat/ConnectFlowSurface", () => ({ default: mockSurface })); describe("ConnectPage", () => { afterEach(() => { - state.oauthReturn = null; - state.connectFlow = null; - state.connectClient = null; - mockReplace.mockClear(); - mockPanel.mockClear(); - mockBanner.mockClear(); + mockSurface.mockClear(); }); - it("renders the MCP connect panel with the user's access token", () => { + it("renders the gateway connect surface with the user's access token and an empty selection", () => { render(); - expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument(); - expect(mockPanel.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] }); - }); - - it("strips the mcpOauthReturn param from the URL after an OAuth return", () => { - state.oauthReturn = "apps"; - window.history.replaceState({}, "", "/connect?mcpOauthReturn=apps"); - render(); - expect(mockReplace).toHaveBeenCalledWith("/connect"); - }); - - it("does not rewrite the URL when there is no OAuth return param", () => { - render(); - expect(mockReplace).not.toHaveBeenCalled(); - }); - - it("mounts the gateway connect banner and puts the panel in connect mode for a DCR flow", () => { - state.connectFlow = "flow-handle-123"; - state.connectClient = "https://claude.ai"; - render(); - expect(screen.getByTestId("connect-flow-banner")).toBeInTheDocument(); - expect(mockBanner.mock.calls[0][0]).toMatchObject({ - flowHandle: "flow-handle-123", - clientOrigin: "https://claude.ai", - }); - expect(mockPanel.mock.calls[0][0].connectMode).toBe(true); - }); - - it("shows no connect banner and leaves connect mode off for a plain visit", () => { - render(); - expect(screen.queryByTestId("connect-flow-banner")).not.toBeInTheDocument(); - expect(mockBanner).not.toHaveBeenCalled(); - expect(mockPanel.mock.calls[0][0].connectMode).toBe(false); - }); - - it("keeps the connect flow handle in the URL while stripping the OAuth return param", () => { - state.oauthReturn = "apps"; - state.connectFlow = "flow-handle-123"; - window.history.replaceState({}, "", "/connect?connect_flow=flow-handle-123&mcpOauthReturn=apps"); - render(); - expect(mockReplace).toHaveBeenCalledWith("/connect?connect_flow=flow-handle-123"); + expect(screen.getByTestId("connect-flow-surface")).toBeInTheDocument(); + expect(mockSurface.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] }); }); }); diff --git a/ui/litellm-dashboard/src/app/connect/page.tsx b/ui/litellm-dashboard/src/app/connect/page.tsx index 3f0c269e86b..652c044197b 100644 --- a/ui/litellm-dashboard/src/app/connect/page.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.tsx @@ -1,36 +1,19 @@ "use client"; -import { Suspense, useEffect, useState } from "react"; -import { useRouter, useSearchParams } from "next/navigation"; +import { Suspense, useState } from "react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; -import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface"; function ConnectPageContent() { const { accessToken } = useAuthorized(); const [selectedServers, setSelectedServers] = useState([]); - const router = useRouter(); - const searchParams = useSearchParams(); - const oauthReturn = searchParams.get("mcpOauthReturn"); - const connectFlow = searchParams.get("connect_flow"); - const connectClient = searchParams.get("connect_client"); - - useEffect(() => { - if (oauthReturn) { - const url = new URL(window.location.href); - url.searchParams.delete("mcpOauthReturn"); - router.replace(url.pathname + url.search); - } - }, [oauthReturn, router]); return (
- {connectFlow && } -
); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx index b3cd6e229af..5caf15d1fce 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx @@ -1,59 +1,61 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { render, screen } from "@testing-library/react"; +import type { ConnectFlowStatus } from "@/components/networking"; import ConnectFlowBanner, { isLoopbackOrigin } from "./ConnectFlowBanner"; vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "https://gateway.example.com", })); +vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ + useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle" }), +})); + afterEach(() => { vi.restoreAllMocks(); - sessionStorage.clear(); }); +const unscoped = (client_origin: string): ConnectFlowStatus => ({ + state: "unscoped", + client_origin, + server_id: null, + server_name: null, + connected: null, +}); + +const renderBanner = (clientOrigin: string) => + render( + , + ); + describe("ConnectFlowBanner", () => { - it("posts the flow handle to the proxy /authorize/complete as a full-page form", () => { - const { container } = render(); + it("posts only the flow handle to the proxy /authorize/complete as a full-page form", () => { + const { container } = renderBanner("https://claude.ai"); const form = container.querySelector("form")!; expect(form).toHaveAttribute("method", "POST"); expect(form).toHaveAttribute("action", "https://gateway.example.com/authorize/complete"); - - const hidden = form.querySelector('input[name="flow"]') as HTMLInputElement; - expect(hidden.value).toBe("flow-handle-123"); - // No token, code, or secret is ever placed in the form; the sealed cookie carries them. + expect(screen.getByDisplayValue("flow-handle-123")).toHaveAttribute("name", "flow"); expect(form.innerHTML).not.toContain("token"); - }); - - it("shows the client origin so the user knows what they are connecting to", () => { - render(); - expect(screen.getAllByText(/claude\.ai/).length).toBeGreaterThan(0); expect(screen.getByRole("button", { name: /finish connecting/i })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument(); }); - it("falls back to a generic label when the client origin is unknown", () => { - render(); - expect(screen.getAllByText(/the application/).length).toBeGreaterThan(0); - }); - - it("offers manual delivery for a loopback client, posted only when checked", () => { - const { container } = render( - , - ); - - const checkbox = container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement; - expect(checkbox).not.toBeNull(); + it("offers manual delivery only for a loopback client, posted only when checked", () => { + const loopback = renderBanner("http://localhost:3118"); + const checkbox = loopback.container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement; expect(checkbox.value).toBe("manual"); expect(checkbox.checked).toBe(false); - expect(screen.getByText(/remote or SSH machine/i)).toBeInTheDocument(); - }); + loopback.unmount(); - it("does not offer manual delivery for a routable client origin or an unknown one", () => { - const routable = render(); + const routable = renderBanner("https://claude.ai"); expect(routable.container.querySelector('input[name="delivery"]')).toBeNull(); - - const unknown = render(); - expect(unknown.container.querySelector('input[name="delivery"]')).toBeNull(); }); it("classifies loopback origins like the server does", () => { @@ -70,12 +72,9 @@ describe("ConnectFlowBanner", () => { }); it("does NOT complete the flow on pagehide (completion requires the explicit button)", () => { - // Security regression: an attacker could lure a signed-in victim to their own client's - // authorize URL; the victim merely closing the tab must NOT deliver a victim-bound code. - // Completion is a deliberate button press, never a side effect of leaving the page. const beaconMock = vi.fn(() => true); vi.stubGlobal("navigator", { ...navigator, sendBeacon: beaconMock }); - render(); + renderBanner("https://claude.ai"); window.dispatchEvent(new Event("pagehide")); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index cea42f916f8..0d6e708f734 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -2,30 +2,18 @@ import React from "react"; import { CheckCircle } from "lucide-react"; -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, ConnectFlowStatus } from "@/components/networking"; +import { OAuth2ConnectButton } from "@/components/chat/MCPAppsPanel"; interface Props { flowHandle: string; - clientOrigin: string | null; + flow?: ConnectFlowStatus; + accessToken: string; + onConnected: () => void; + failed: boolean; } -/** - * The interlude shown when a DCR client (Claude Desktop, MCP Inspector) sends the user - * through the gateway sign-in and lands them on the apps grid to authorize servers. The - * grid below authorizes individual servers into the per-user vault; this banner is the - * finish step that returns the user to the client. - * - * Finishing requires the explicit "Finish connecting" button: a native form POST to the proxy's - * /authorize/complete, which mints the gateway authorization code and 303-redirects to the DCR - * client's own redirect URI (the full-page navigation carries the HttpOnly per-flow cookie and - * follows the cross-origin redirect to the client's loopback). - * - * The button press IS the consent gate and must not be bypassed. An earlier version auto-finished - * on tab close via navigator.sendBeacon; that let an attacker who lured a signed-in victim to their - * own client's authorize URL harvest a victim-bound code the moment the victim closed the tab - * (no click). Merely visiting the authorize URL is attacker-inducible, so completion has to be a - * deliberate user action, not a side effect of leaving the page. - */ +/** Finish remains an explicit POST because a cross-site navigation must never mint a code. */ export function isLoopbackOrigin(origin: string | null): boolean { if (!origin) return false; try { @@ -36,10 +24,44 @@ export function isLoopbackOrigin(origin: string | null): boolean { } } -const ConnectFlowBanner: React.FC = ({ flowHandle, clientOrigin }) => { +const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly [string, string] => { + const clientLabel = flow?.client_origin ?? "the application"; + const serverLabel = flow?.server_name ?? "the requested MCP server"; + if (failed || flow === undefined || flow.state === "stale") { + return [ + "The connection cannot continue", + `The gateway could not validate this connection. Cancel to return to ${clientLabel}.`, + ]; + } + if (flow.state === "unscoped") { + return [ + `Connect your MCP servers to ${clientLabel}`, + `Authorize the servers you want to use below, then click Finish connecting to return to ${clientLabel}.`, + ]; + } + if (flow.state === "interactive" && !flow.connected) { + return [ + `Allow ${clientLabel} to use ${serverLabel}`, + `Authorize ${serverLabel} below to continue, or cancel to send ${clientLabel} away.`, + ]; + } + return [ + `Allow ${clientLabel} to use ${serverLabel}`, + `Click Finish connecting to give ${clientLabel} access to ${serverLabel} as you.`, + ]; +}; + +const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onConnected, failed }) => { const action = `${getProxyBaseUrl()}/authorize/complete`; - const clientLabel = clientOrigin ?? "the application"; - const loopbackClient = isLoopbackOrigin(clientOrigin); + const state = failed || flow === undefined ? "stale" : flow.state; + const canFinish = state === "unscoped" || (state !== "stale" && flow?.connected === true); + const canCancel = state !== "unscoped"; + const loopbackClient = isLoopbackOrigin(flow?.client_origin ?? null); + const vendorServer = + state === "interactive" && flow?.connected === false && flow.server_id !== null + ? { server_id: flow.server_id, server_name: flow.server_name } + : null; + const copy = copyFor(flow, failed); return (
@@ -47,27 +69,48 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, clientOrigin }) => {
-

Connect your MCP servers to {clientLabel}

-

- Authorize the servers you want to use below, then click Finish connecting to return to {clientLabel}. -

+

{copy[0]}

+

{copy[1]}

-
- - - {loopbackClient && ( - +
+ {vendorServer !== null && ( + )} - +
+ + {canFinish && ( + + )} + {canCancel && ( + + )} + {loopbackClient && ( + + )} +
+
); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx new file mode 100644 index 00000000000..cf59dcfc368 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx @@ -0,0 +1,118 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { act, render, screen, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import ConnectFlowSurface from "./ConnectFlowSurface"; +import { fetchConnectFlow } from "@/components/networking"; + +const { startOAuthFlow, state, onSuccess } = vi.hoisted(() => ({ + startOAuthFlow: vi.fn(), + onSuccess: { current: undefined as (() => void) | undefined }, + state: { oauthReturn: null as string | null, connectFlow: null as string | null }, +})); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ replace: vi.fn() }), + useSearchParams: () => ({ + get: (key: string) => ({ mcpOauthReturn: state.oauthReturn, connect_flow: state.connectFlow })[key] ?? null, + }), +})); +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + fetchConnectFlow: vi.fn(), + getProxyBaseUrl: () => "https://gateway.example.com", +})); +vi.mock("@/components/chat/MCPAppsPanel", async (importOriginal) => ({ + ...(await importOriginal()), + default: () =>
, +})); +vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ + useUserMcpOAuthFlow: ({ onSuccess: success }: { onSuccess: () => void }) => { + onSuccess.current = success; + return { startOAuthFlow, status: "idle" }; + }, +})); + +const flow = (state: "unscoped" | "interactive" | "m2m" | "stale", connected: boolean | null = null) => ({ + state, + client_origin: "https://claude.ai", + server_id: state === "interactive" || state === "m2m" ? "s-design" : null, + server_name: state === "interactive" || state === "m2m" ? "design_tool" : null, + connected, +}); + +const renderSurface = () => + render( + + + , + ); + +afterEach(() => { + state.oauthReturn = null; + state.connectFlow = null; + onSuccess.current = undefined; + sessionStorage.clear(); + vi.clearAllMocks(); +}); + +describe("ConnectFlowSurface", () => { + it.each([ + { result: flow("unscoped"), grid: true, finish: true, cancel: false, oauthStarts: 0 }, + { result: flow("interactive", false), grid: false, finish: false, cancel: true, oauthStarts: 1 }, + { result: flow("interactive", true), grid: false, finish: true, cancel: true, oauthStarts: 0 }, + { result: flow("m2m", true), grid: false, finish: true, cancel: true, oauthStarts: 0 }, + { result: flow("stale"), grid: false, finish: false, cancel: true, oauthStarts: 0 }, + ])( + "renders $result.state without widening its action surface", + async ({ result, grid, finish, cancel, oauthStarts }) => { + state.connectFlow = "flow-handle-123"; + vi.mocked(fetchConnectFlow).mockResolvedValue(result); + renderSurface(); + + await screen.findByRole("button", { name: /finish connecting|cancel|connect/i }); + await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledTimes(oauthStarts)); + expect(screen.queryByTestId("mcp-apps-panel") !== null).toBe(grid); + expect(screen.queryByRole("button", { name: /finish connecting/i }) !== null).toBe(finish); + expect(screen.queryByRole("button", { name: "Cancel" }) !== null).toBe(cancel); + }, + ); + + it("keeps the grid and Finish hidden until the gateway accepts a handle", () => { + state.connectFlow = "flow-handle-123"; + vi.mocked(fetchConnectFlow).mockReturnValue(new Promise(() => {})); + renderSurface(); + + expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Cancel" })).toHaveAttribute("value", "deny"); + }); + + it("keeps the grid and Finish hidden when flow validation fails", async () => { + state.connectFlow = "invalid-handle"; + vi.mocked(fetchConnectFlow).mockRejectedValue(new Error("invalid flow")); + renderSurface(); + + await screen.findByRole("button", { name: "Cancel" }); + expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); + }); + + it("refetches the sealed flow after the vendor connection completes", async () => { + state.connectFlow = "flow-handle-123"; + vi.mocked(fetchConnectFlow) + .mockResolvedValueOnce(flow("interactive", false)) + .mockResolvedValueOnce(flow("interactive", true)); + renderSurface(); + + await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledOnce()); + await act(async () => onSuccess.current?.()); + + await screen.findByRole("button", { name: /finish connecting/i }); + }); + + it("renders the ordinary panel without a flow handle", () => { + renderSurface(); + expect(fetchConnectFlow).not.toHaveBeenCalled(); + expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx new file mode 100644 index 00000000000..33a61fceacf --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx @@ -0,0 +1,59 @@ +"use client"; + +import React, { useEffect } from "react"; +import { useRouter, useSearchParams } from "next/navigation"; +import { useQuery } from "@tanstack/react-query"; +import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; +import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import { fetchConnectFlow } from "@/components/networking"; + +interface Props { + accessToken: string; + selectedServers: string[]; + onChange: (servers: string[]) => void; +} + +/** Renders the sealed gateway connect flow without trusting URL context. */ +const ConnectFlowSurface: React.FC = ({ accessToken, selectedServers, onChange }) => { + const router = useRouter(); + const searchParams = useSearchParams(); + const oauthReturn = searchParams.get("mcpOauthReturn"); + const connectFlow = searchParams.get("connect_flow"); + + useEffect(() => { + if (oauthReturn) { + const url = new URL(window.location.href); + url.searchParams.delete("mcpOauthReturn"); + router.replace(url.pathname + url.search); + } + }, [oauthReturn, router]); + + const flowQuery = { + queryKey: ["gateway-connect-flow", connectFlow], + queryFn: () => fetchConnectFlow(connectFlow!), + enabled: !!connectFlow, + retry: false, + }; + const { data: flow, isError, refetch } = useQuery(flowQuery); + + if (connectFlow === null) { + return ; + } + + return ( + <> + + {flow?.state === "unscoped" && ( + + )} + + ); +}; + +export default ConnectFlowSurface; diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx index e8795c36bc6..c9609405676 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx @@ -88,6 +88,13 @@ describe("MCPAppsPanel logos", () => { }); const connectServers = [ + { + server_id: "s-m2m", + server_name: "service_tool", + auth_type: "oauth2", + oauth2_flow: "client_credentials", + connected_app_reachable: true, + }, { server_id: "s-reach", server_name: "reachable_srv", @@ -124,6 +131,8 @@ describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => { expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true); expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument(); expect(screen.getByText("Connected (1)")).toBeInTheDocument(); + expect(screen.getByText("service_tool")).toBeInTheDocument(); + expect(screen.queryByText("Connect", { exact: true })).not.toBeInTheDocument(); const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]); expect(toolCountFetchedIds).toContain("s-reach"); expect(toolCountFetchedIds).not.toContain("s-unreach"); diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 38a095c067d..1fee8923e94 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -13,23 +13,32 @@ import { getMCPOAuthUserCredentialStatus, listMCPTools, } from "../networking"; -import { AUTH_TYPE, MCPServer, MCPTool, handleTransport, isUnsupportedOnGatewayConnect } from "../mcp_tools/types"; +import { + getMcpOAuthMode, + MCPServer, + MCPTool, + handleTransport, + isUnsupportedOnGatewayConnect, +} from "../mcp_tools/types"; import { Logo } from "@/components/molecules/logo/Logo"; import { toast } from "@/lib/toast"; import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; interface OAuth2ConnectButtonProps { - server: MCPServer; + server: Pick; accessToken: string; onConnect: (serverId: string) => void; variant?: "badge" | "button"; + autoStartKey?: string | null; } -const OAuth2ConnectButton: React.FC = ({ +export const OAuth2ConnectButton: React.FC = ({ server, accessToken, onConnect, variant = "badge", + autoStartKey = null, }) => { const name = server.server_name ?? server.alias ?? server.server_id; const { startOAuthFlow, status } = useUserMcpOAuthFlow({ @@ -39,6 +48,12 @@ const OAuth2ConnectButton: React.FC = ({ onSuccess: useCallback(() => onConnect(server.server_id), [onConnect, server.server_id]), }); + useEffect(() => { + if (autoStartKey === null || status !== "idle" || getSecureItem(autoStartKey) !== null) return; + setSecureItem(autoStartKey, "1"); + startOAuthFlow(); + }, [autoStartKey, status, startOAuthFlow]); + const loading = status === "authorizing" || status === "exchanging"; if (variant === "button") { @@ -190,7 +205,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, if (!isCurrentLoad()) return; const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? []; const reachable = connectMode ? list.filter((s) => s.connected_app_reachable !== false) : list; - const oauthServers = reachable.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2); + const oauthServers = reachable.filter((s) => getMcpOAuthMode(s) === "authorization_code"); commitServers(reachable); setOauthChecking(new Set(oauthServers.map((s) => s.server_id))); setLoading(false); @@ -274,7 +289,10 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, {unavailabilityLabel} ); } - if (server.auth_type === AUTH_TYPE.OAUTH2) { + if (getMcpOAuthMode(server) === "m2m") { + return ; + } + if (getMcpOAuthMode(server) === "authorization_code") { if (oauthConnected.has(server.server_id)) { return ; } @@ -339,7 +357,10 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, if (unavailabilityLabel !== null) { return {unavailabilityLabel}; } - if (detailServer.auth_type !== AUTH_TYPE.OAUTH2) { + if (getMcpOAuthMode(detailServer) === "m2m") { + return Authorized; + } + if (getMcpOAuthMode(detailServer) !== "authorization_code") { return (