diff --git a/.github/workflows/guard-main-branch.yml b/.github/workflows/guard-main-branch.yml deleted file mode 100644 index 5bc561c6441..00000000000 --- a/.github/workflows/guard-main-branch.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: Guard main branch - -on: - pull_request: - branches: - - main - merge_group: - -permissions: {} - -# DO NOT RENAME the job's `name:` — it is referenced by GitHub branch -# protection as a required status check on `main`. Renaming silently -# breaks the gate. -jobs: - guard: - name: Verify PR source branch - runs-on: ubuntu-latest - timeout-minutes: 2 - steps: - - name: Reject merge_group events - if: github.event_name == 'merge_group' - run: | - echo "::error::Merge queue is not supported for main. Disable merge queue or update this guard." - exit 1 - - name: Check head branch name - env: - HEAD_REF: ${{ github.head_ref }} - HEAD_REPO: ${{ github.event.pull_request.head.repo.full_name }} - BASE_REPO: ${{ github.repository }} - run: | - echo "PR head repo: $HEAD_REPO" - echo "PR head branch: $HEAD_REF" - if [ "$HEAD_REPO" != "$BASE_REPO" ]; then - echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against 'litellm_internal_staging' instead." - exit 1 - fi - if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then - echo "Allowed source branch." - exit 0 - fi - echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_internal_staging' instead." - exit 1 diff --git a/litellm/constants.py b/litellm/constants.py index 6b984c2673c..be5a4e4b076 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -571,6 +571,7 @@ ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int( LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) +LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS: Final = 5.0 LOGGING_WORKER_CLEAR_PERCENTAGE: Final = int( os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50) ) # Percentage of queue to clear (default: 50%) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 1d74595781a..5ccc5632646 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -18,11 +18,16 @@ from litellm.constants import ( LOGGING_WORKER_CONCURRENCY, LOGGING_WORKER_MAX_QUEUE_SIZE, LOGGING_WORKER_MAX_TIME_PER_COROUTINE, + LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS, MAX_ITERATIONS_TO_CLEAR_QUEUE, MAX_TIME_TO_CLEAR_QUEUE, ) +def _coroutine_name(coroutine: Coroutine) -> str: + return getattr(coroutine, "__qualname__", None) or getattr(coroutine, "__name__", None) or type(coroutine).__name__ + + class LoggingTask(TypedDict): """ A logging task with its associated context to ensure logging is executed in @@ -47,10 +52,12 @@ class LoggingWorker: timeout: float = LOGGING_WORKER_MAX_TIME_PER_COROUTINE, max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE, concurrency: int = LOGGING_WORKER_CONCURRENCY, + timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS, ): self.timeout = timeout self.max_queue_size = max_queue_size self.concurrency = concurrency + self.timeout_summary_window = timeout_summary_window self._queue: asyncio.Queue[LoggingTask] | None = None self._worker_task: asyncio.Task | None = None self._running_tasks: set[asyncio.Task] = set() @@ -59,6 +66,10 @@ class LoggingWorker: self._bound_loop: asyncio.AbstractEventLoop | None = None self._last_aggressive_clear_time: float = 0.0 self._aggressive_clear_in_progress: bool = False + self._timeout_total: int = 0 + self._timeout_burst_count: int = 0 + self._timeout_last_callback: str | None = None + self._timeout_summary_task: asyncio.Task | None = None # Register cleanup handler to flush remaining events on exit atexit.register(self._flush_on_exit) @@ -136,6 +147,8 @@ class LoggingWorker: self._sem = None self._worker_task = None self._running_tasks.clear() + self._timeout_summary_task = None + self._timeout_burst_count = 0 self._queue = new_queue self._bound_loop = current_loop return @@ -156,12 +169,15 @@ class LoggingWorker: """Runs the logging task and handles cleanup. Releases semaphore when done.""" try: if self._queue is not None: + # Run the coroutine in its original context + callback_task: Final = task["context"].run(asyncio.create_task, task["coroutine"]) try: - # Run the coroutine in its original context - await asyncio.wait_for( - task["context"].run(asyncio.create_task, task["coroutine"]), - timeout=self.timeout, - ) + await asyncio.wait_for(callback_task, timeout=self.timeout) + except asyncio.TimeoutError as e: + if callback_task.cancelled(): + self._record_callback_timeout(task["coroutine"]) + else: + verbose_logger.exception("LoggingWorker error: %s", e) except Exception as e: verbose_logger.exception("LoggingWorker error: %s", e) finally: @@ -171,6 +187,35 @@ class LoggingWorker: # Always release semaphore, even if queue is None sem.release() + def _record_callback_timeout(self, coroutine: Coroutine) -> None: + """Count a callback timeout and arm a debounced summary, so a burst of timeouts + (e.g. a slow Redis timing out many callbacks at once) logs one bounded line rather + than a full ERROR stacktrace per callback.""" + self._timeout_total += 1 + self._timeout_burst_count += 1 + self._timeout_last_callback = _coroutine_name(coroutine) + if self._timeout_summary_task is None or self._timeout_summary_task.done(): + self._timeout_summary_task = asyncio.create_task(self._flush_timeout_summary()) + + async def _flush_timeout_summary(self) -> None: + """After the burst settles, log one bounded summary covering every timeout in it.""" + await asyncio.sleep(self.timeout_summary_window) + self._emit_timeout_summary() + + def _emit_timeout_summary(self) -> None: + """Log one bounded summary for the current burst and reset the burst counter.""" + burst_count: Final = self._timeout_burst_count + self._timeout_burst_count = 0 + if burst_count <= 0: + return + verbose_logger.warning( + "LoggingWorker: %d callback(s) timed out after %ss (callback: %s); %d timed out since start", + burst_count, + self.timeout, + self._timeout_last_callback, + self._timeout_total, + ) + async def _worker_loop(self) -> None: """Main worker loop that gets tasks and schedules them to run concurrently.""" try: @@ -406,6 +451,11 @@ class LoggingWorker: async def stop(self) -> None: """Stop the logging worker and clean up resources.""" + if self._timeout_summary_task is not None: + self._timeout_summary_task.cancel() + self._timeout_summary_task = None + self._emit_timeout_summary() + if self._worker_task is None and not self._running_tasks: # No worker launched and no in-flight tasks to drain. return diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index c5290b41f7b..a446a721307 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -50,6 +50,14 @@ class BaseLLMModelInfo(ABC): """ return None + def get_model_cost_key(self, model: str) -> str | None: + """ + Maps the model name a user sends to the key `litellm.model_cost` stores it under, when the two differ. + `get_model_info` tries this key once the exact `model` and `provider/model` keys miss. The default None means + the provider's user-facing names already match the cost map, so there is nothing extra to try. + """ + return None + @abstractmethod def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: """ diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index b4e1856f499..05160d83c12 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -602,6 +602,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return None return max(matches, key=lambda match: len(match[0]))[1] + def get_model_cost_key(self, model: str) -> str: + return f"fireworks_ai/{resolve_fireworks_resource_name(model)}" + def get_provider_info(self, model: str) -> ProviderSpecificModelInfo: supports_function_calling_value: Final = self._get_model_cost_capability( model=model, capability="supports_function_calling" diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index 6e1d121c3be..a35e5d1ce10 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -67,9 +67,8 @@ module materially harder to understand. auth, SSE, streamable HTTP, and stdio as separate flows. Do not collapse them behind a single generic branch unless tests prove every mode still behaves correctly. -- Be especially careful with `available_on_public_internet: false` combined with - `delegate_auth_to_upstream: true`. The local `CLAUDE.md` explains the anonymous - upstream PKCE path that must remain intentional. +- Be especially careful with legacy `delegate_auth_to_upstream: true`. The local + `CLAUDE.md` explains its admitted replacement and public discovery contract. - Keep database-backed fields in sync across migrations, typed models under `litellm/types/mcp.py` or `litellm/types/mcp_server/`, config loading, this package, and dashboard state when the field is user-visible. diff --git a/litellm/proxy/_experimental/mcp_server/CLAUDE.md b/litellm/proxy/_experimental/mcp_server/CLAUDE.md index 0ba8f73315f..7f8d06b4570 100644 --- a/litellm/proxy/_experimental/mcp_server/CLAUDE.md +++ b/litellm/proxy/_experimental/mcp_server/CLAUDE.md @@ -1 +1 @@ -MCP note: **`available_on_public_internet: false` with `delegate_auth_to_upstream: true` (oauth2, interactive - not `client_credentials`)** - LiteLLM still allows the anonymous upstream PKCE path (no proxy API key for `/authorize` and matching MCP routes). The internal-only flag mainly affects other surfaces (e.g. IP-based discovery). Rely on the upstream IdP and network policy; the dashboard shows a warning when both are set, and the proxy logs a warning when the server is loaded from config or the database +MCP note: **`auth_type: oauth2` with `delegate_auth_to_upstream: true` is deprecated** - LiteLLM admission is required for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth. OAuth discovery endpoints stay public so clients can start the RFC 9728 flow 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 cd6739777dd..b0d57cb6228 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 @@ -129,10 +129,9 @@ def _is_mcp_passthrough_cold_start(mcp_servers: list[str] | None, client_ip: str spec-compliant WWW-Authenticate challenge instead of surfacing a generic admission error. - Uses "all" semantics (mirrors - :meth:`MCPRequestHandler._target_servers_delegate_auth_to_upstream`): one - non-passthrough target in a co-targeted set must not flip the bypass open - for the others. Fails closed when any target cannot be resolved.""" + Uses "all" semantics: one non-passthrough target in a co-targeted set must + not flip the bypass open for the others. Fails closed when any target + cannot be resolved.""" if not mcp_servers: return False from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -146,6 +145,27 @@ def _is_mcp_passthrough_cold_start(mcp_servers: list[str] | None, client_ip: str return True +def _is_legacy_delegate_cold_start(mcp_servers: list[str] | None, client_ip: str | None) -> bool: + """Allow only credential-free legacy delegates to reach the route's OAuth challenge.""" + if not mcp_servers: + return False + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + + for name in mcp_servers: + server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip) + if server is None or server.auth_type != MCPAuth.oauth2: + return False + if server.delegate_auth_to_upstream is not True: + return False + if MCPServerManager.effective_oauth2_flow(server) == "client_credentials": + return False + return True + + def _is_litellm_auth_admission_error(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code == 401 @@ -277,9 +297,18 @@ def _admission_failure_fallback( mcp_servers_from_path is not None and not _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers) and _is_litellm_auth_admission_error(exc) - and _is_mcp_passthrough_cold_start( - mcp_servers_from_path, - client_ip=IPAddressUtils.get_mcp_client_ip(request), + and ( + _is_mcp_passthrough_cold_start( + mcp_servers_from_path, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ) + or ( + not bearer_presented + and _is_legacy_delegate_cold_start( + mcp_servers_from_path, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ) + ) ) ): verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") @@ -434,22 +463,6 @@ class MCPRequestHandler: api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}", request=request, ) - elif MCPRequestHandler._target_servers_delegate_auth_to_upstream( - path=request_route, - mcp_servers=mcp_servers, - client_ip=IPAddressUtils.get_mcp_client_ip(request), - ): - # Operator opted this oauth2 server into upstream-delegated auth: the - # client authenticates directly with the upstream MCP server, so any - # Authorization bearer is an upstream token, never a LiteLLM key. Skip - # LiteLLM validation entirely — covering both the no-credential - # discovery request and the authenticated call carrying the upstream - # bearer — so a tool call that succeeds never carries a phantom 401 - # auth span; the bearer is forwarded upstream unchanged. Gated by - # _target_servers_delegate_auth_to_upstream, which returns True only - # when EVERY target is auth_type=oauth2 with delegate_auth_to_upstream - # set; fails closed otherwise. - validated_user_api_key_auth = UserAPIKeyAuth() elif MCPRequestHandler._target_servers_are_true_passthrough( path=request_route, mcp_servers=mcp_servers, @@ -660,64 +673,6 @@ class MCPRequestHandler: return [single_server_match.group(1)] return [servers_and_path] - @staticmethod - def _target_servers_delegate_auth_to_upstream( - path: str, mcp_servers: list[str] | None, client_ip: str | None - ) -> bool: - """ - True only when EVERY MCP server the request targets is configured for - ``auth_type == oauth2`` AND has ``delegate_auth_to_upstream=True``. - Fails closed when any target does not opt in or cannot be resolved. - - Used by :meth:`process_mcp_request` to skip LiteLLM API-key/SSO auth - entirely (PKCE passthrough) so the client authenticates directly with - the upstream MCP server. Mixed-target requests (e.g. one delegated + - one non-delegated server) fall back to normal LiteLLM auth. - """ - # Inline imports avoid a circular dependency: mcp_server_manager imports - # from this module. - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - - # Must mirror the downstream header-vs-path override - # (``extract_mcp_auth_context``) or an attacker could set - # ``x-mcp-servers`` to a delegate-enabled server while the URL path - # targets a non-delegate server, skipping LiteLLM auth for it. - target_names: Final = MCPRequestHandler._resolve_target_server_names(path=path, mcp_servers_header=mcp_servers) - if not target_names: - return False - - for name in target_names: - server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip) - if server is None or server.auth_type != MCPAuth.oauth2: - return False - # `is True` is intentional: opt-in must be an explicit boolean - # True. A MagicMock attribute (in tests) or any other truthy - # non-bool must not silently enable the bypass. - if getattr(server, "delegate_auth_to_upstream", False) is not True: - return False - # Never delegate for M2M (client_credentials) servers: LiteLLM - # fetches the upstream token automatically using stored credentials, - # so allowing anonymous bypass would let any external caller invoke - # tools authenticated as LiteLLM's service account. - # - # Resolve the flow rather than reading has_client_credentials directly: - # this is a security gate, and a legacy row whose oauth2_flow was never - # stamped still carries the M2M credential shape (client_id/secret + - # token_url, no authorization_url). Treating an unstamped-but-M2M-shaped - # row as non-M2M here would reopen the anonymous bypass the explicit - # column no longer closes on its own. Shares the one resolution helper - # with the egress backstop and the anonymous-delegate allowlist; all fail - # closed on the ambiguous shape and are removed together once no null rows - # remain. A pure-PKCE delegate server (no stored credentials) resolves to a - # non-M2M flow and keeps its bypass. - if MCPServerManager.effective_oauth2_flow(server) == "client_credentials": - return False - return True - @staticmethod def _target_servers_are_true_passthrough(path: str, mcp_servers: list[str] | None, client_ip: str | None) -> bool: """ @@ -726,7 +681,7 @@ class MCPRequestHandler: Used by :meth:`process_mcp_request` to skip LiteLLM admission auth entirely: the gateway is a transparent proxy and the caller's ``Authorization`` is an upstream token, never a LiteLLM key. - Mirrors :meth:`_target_servers_delegate_auth_to_upstream`; a mixed-target request keeps normal auth. + A mixed-target request keeps normal auth. """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3291d5effc2..fb0c623473a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1055,7 +1055,7 @@ def _should_strip_caller_authorization( pass-through cold-start case (RFC 9728) the bearer in ``Authorization`` is the upstream OAuth token and must be forwarded, so we keep it. - - **oauth_delegate servers**: admission always runs and there is no + - **Delegated OAuth servers**: admission always runs and there is no anonymous path, so the caller's separate ``Authorization`` is forwarded only when a distinct ``x-litellm-api-key`` carried admission. Without that header the ``Authorization`` *was* the @@ -1075,12 +1075,17 @@ def _should_strip_caller_authorization( # upstream — it would override another user's stored credential. Delegate and # pass-through return None from to_server_spec and keep forwarding the bearer. return True - if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate): + is_delegated_oauth: Final = mcp_server.is_oauth_delegate or ( + mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream + ) + if not (mcp_server.is_oauth_passthrough or is_delegated_oauth): return False has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers) - if mcp_server.is_oauth_delegate: - return not has_explicit_litellm_admission_header + if is_delegated_oauth: + return not has_explicit_litellm_admission_header or _authorization_is_litellm_admission_credential( + raw_headers, user_api_key_auth + ) return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or ( user_api_key_auth is None and not has_explicit_litellm_admission_header ) @@ -1107,15 +1112,11 @@ def _authorization_is_litellm_admission_credential( That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the same key in both headers. """ - if user_api_key_auth is None or not user_api_key_auth.api_key: - return False admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key") - if not admission_header: - return True authorization: Final = _raw_header_value(raw_headers, "authorization") - return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme( - admission_header, "Bearer" - ) + if admission_header and authorization: + return strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(admission_header, "Bearer") + return bool(user_api_key_auth and user_api_key_auth.api_key and not admission_header) def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str: @@ -1453,22 +1454,19 @@ def _warn_on_server_name_fields( _warn("server_name", server_name) -def _warn_internal_delegate_pkce_if_applicable(server: MCPServer, *, source: str) -> None: - """Surface internal + upstream PKCE delegate in logs for operators.""" +def _warn_legacy_delegate_auth_if_applicable(server: MCPServer, *, source: str) -> None: + """Direct legacy delegated OAuth configurations to the admitted replacement.""" if server.auth_type != MCPAuth.oauth2: return if getattr(server, "delegate_auth_to_upstream", False) is not True: return - if getattr(server, "available_on_public_internet", True): - return if server.has_client_credentials: return label: Final = get_server_prefix(server) verbose_logger.warning( - "MCP server %r (id=%s, source=%s): internal-only (available_on_public_internet=false) " - "with delegate_auth_to_upstream=true. Anonymous callers can reach the upstream OAuth2 " - "/authorize flow and complete PKCE without a LiteLLM API key session; ensure the " - "upstream IdP and network enforce your access policy.", + "MCP server %r (id=%s, source=%s) uses deprecated auth_type=oauth2 with " + "delegate_auth_to_upstream=true. LiteLLM admission is now required; migrate to " + "auth_type=oauth_delegate for client-forwarded OAuth.", label, server.server_id, source, @@ -2640,7 +2638,7 @@ class MCPServerManager: oauth_identity_binding=server_config.get("oauth_identity_binding", None), ) self._assign_unique_short_prefix(new_server) - _warn_internal_delegate_pkce_if_applicable(new_server, source="config") + _warn_legacy_delegate_auth_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) self._invalidate_discovery_lists(server_id) self.config_mcp_servers[server_id] = new_server @@ -3185,7 +3183,7 @@ class MCPServerManager: timeout=getattr(mcp_server, "timeout", None), max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None), ) - _warn_internal_delegate_pkce_if_applicable(new_server, source="database") + _warn_legacy_delegate_auth_if_applicable(new_server, source="database") self._set_oauth_discovery_deferred( new_server.server_id, _requires_oauth_discovery(server_url, use_issuer_anchor, new_server), @@ -3479,10 +3477,6 @@ class MCPServerManager: ) ) - # For anonymous callers (no user_id, no role), also surface any - # servers the operator has opted into upstream-delegated auth. - # These servers handle their own auth at the upstream level, so - # LiteLLM granting access here does not bypass any security gate. is_anonymous: Final = not ( user_api_key_auth and ( @@ -3492,23 +3486,12 @@ class MCPServerManager: ) ) if is_anonymous: - delegate_server_ids: Final = [ + passthrough_server_ids: Final = [ server.server_id for server in self.get_registry().values() - if ( - getattr(server, "auth_type", None) == MCPAuth.oauth2 - and getattr(server, "delegate_auth_to_upstream", False) is True - # M2M servers must not be exposed anonymously: an - # unauthenticated caller would get LiteLLM to proxy tool - # calls using its stored client_credentials. Resolve the flow - # rather than reading has_client_credentials so an unstamped - # M2M-shape row (null column, verbatim-read as non-M2M) still - # fails closed here, matching the anonymous-delegate auth gate. - and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" - ) - or getattr(server, "auth_type", None) == MCPAuth.true_passthrough + if getattr(server, "auth_type", None) == MCPAuth.true_passthrough ] - combined_servers.update(delegate_server_ids) + combined_servers.update(passthrough_server_ids) restrict_allow_all: Final = ( resolved_general_settings.get("mcp_allow_all_keys_respects_mcp_scope", False) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index fc87db69e16..7feb1fd468d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -4257,20 +4257,6 @@ if MCP_AVAILABLE: return None return _get_authorization_header_from_scope(scope) - def _is_delegate_upstream_probe_target(server: MCPServer) -> bool: - """Whether ``server`` is an interactive delegate-auth server whose client-supplied - token should be preflighted upstream. - - Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is - resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed - (its stored client credentials drive egress; the caller's bearer is irrelevant). - """ - return ( - server.auth_type == MCPAuth.oauth2 - and server.delegate_auth_to_upstream is True - and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" - ) - async def _probe_upstream_auth( url: str, auth_header: str, @@ -4331,7 +4317,7 @@ if MCP_AVAILABLE: mcp_servers: list[str] | None, client_ip: str | None, ) -> None: - """Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts. + """Probe pass-through upstream servers in parallel before the MCP session starts. Only servers the caller's key is already authorized to reach are probed — the list is derived from _get_allowed_mcp_servers so that a user cannot @@ -4343,38 +4329,9 @@ if MCP_AVAILABLE: if the upstream accepts it but forbids the caller. Fails-open: network errors are logged and the request is allowed through. - Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``) - are probed with the caller's bare ``Authorization`` bearer. That bearer is only - an upstream token (never a LiteLLM key) when admission took the delegate bypass, - so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same - resolver admission used -- rather than the wider allowed-server prefix/access-group - matching. A name that only reaches a delegate server via server_id or an access - group would have been admitted as a real LiteLLM key, so probing it would leak that - key upstream; requiring the admission-resolver match closes that gap. Without the - probe a rejected token is absorbed by the tools/list handler and masked as an empty - tool list. Gated to single-server routes so one rejected token cannot 401 a - multi-server aggregate connect, matching the OBO preflight gating; the challenge - echoes the requested name so aliased routes get the same resource_metadata URL as - the tokenless preemptive challenge. """ forwarded_auth: Final = _get_forwarded_auth_from_scope(scope) - requested_single_target: Final = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None - # The bare Authorization header (no x-litellm-api-key) is a valid upstream token - # only when admission classified it as one, i.e. the single requested name resolves - # to a delegate server under admission's own resolver. Resolve it the same way here - # so a server_id- or access-group-named delegate (which admission would have treated - # as a LiteLLM key) is never probed with that key. - delegate_server: Final = ( - global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip) - if requested_single_target - else None - ) - delegate_auth: Final = ( - _get_authorization_header_from_scope(scope) - if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server) - else None - ) - if not forwarded_auth and not delegate_auth: + if not forwarded_auth: return # Use the authorized server set, not the raw user-supplied names, so that @@ -4384,35 +4341,20 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) - passthrough_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = ( - tuple( - (srv, forwarded_auth, srv.name) - for srv in allowed_servers - # Restrict to genuine OAuth pass-through servers (auth_type none + - # Authorization in extra_headers). Gateway-managed OAuth2 servers - # must not receive the ``resource_metadata=`` challenge emitted - # below — they require ``authorization_uri=`` pointing at the - # gateway AS metadata. ``is_oauth_passthrough`` already requires - # ``auth_type in (None, MCPAuth.none)``, which is mutually - # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), - # so M2M servers are implicitly excluded here. - if srv.is_oauth_passthrough - ) - if forwarded_auth - else () + passthrough_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple( + (srv, forwarded_auth, srv.name) + for srv in allowed_servers + # Restrict to genuine OAuth pass-through servers (auth_type none + + # Authorization in extra_headers). Gateway-managed OAuth2 servers + # must not receive the ``resource_metadata=`` challenge emitted + # below — they require ``authorization_uri=`` pointing at the + # gateway AS metadata. ``is_oauth_passthrough`` already requires + # ``auth_type in (None, MCPAuth.none)``, which is mutually + # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), + # so M2M servers are implicitly excluded here. + if srv.is_oauth_passthrough ) - # Probe the admission-resolved delegate server only when the caller is actually - # authorized for it (present in the IP-filtered allowed set), keyed by server_id. - delegate_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = ( - tuple( - (srv, delegate_auth, requested_single_target) - for srv in allowed_servers - if delegate_server is not None and srv.server_id == delegate_server.server_id - ) - if delegate_auth and requested_single_target - else () - ) - probe_targets: Final = passthrough_targets + delegate_targets + probe_targets: Final = passthrough_targets if not probe_targets: return diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 93bbc567430..c01fe15fc09 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -5,8 +5,10 @@ At request time the spend writer builds one ToolUsageTransaction per request tha invoked tools (MCP namespaced tool name plus response tool_calls; declared-but-not- invoked tools are excluded) and queues it on the prisma client. The spend-log flush job drains the queue into LiteLLM_SpendLogToolIndex (per-request drill-down) and -LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads) in a -single transaction, so a failed flush never leaves a partial rollup increment. +LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads). The +index rows are keyed on (request_id, tool_name) and written with skip_duplicates, +so they go out as bounded standalone statements; every rollup upsert stays in one +transaction, so a failed flush never leaves a partial rollup increment. """ from __future__ import annotations @@ -19,7 +21,10 @@ from datetime import datetime, timezone from itertools import groupby from typing import TYPE_CHECKING, Any, Final +from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES +from litellm.proxy.db.spend_log_batching import spend_log_write_batches +from litellm.repositories.table_repositories import SpendLogToolIndexRepository if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -98,14 +103,19 @@ async def flush_tool_usage_transactions( transactions: Sequence[ToolUsageTransaction], n_retry_times: int = 3, ) -> None: - """Write index rows and rollup upserts for a drained queue batch in one - transaction. Retries only ConnectError, the one failure that proves the - statements never reached the database. Post-send failures (Read timeouts - and errors) are ambiguous and are NOT retried: the engine can abandon the - transaction open on the pooled connection, so a retry's statements stack - into the same transaction and one commit applies both increment sets. - Ambiguous failures drop the batch; the caller logs it at error. Callers - must not add their own retry around this function.""" + """Write the index rows as bounded standalone statements, then every rollup + upsert for the drained queue batch in one transaction. One flush fans out to + transactions x tools index rows, so the index write is split by the spend-log + statement budgets; a split inside ``batch_()`` would not help, since the + batcher ships every queued statement to the query engine as one payload. + Retries only ConnectError, the one failure that proves the statements never + reached the database; replayed index rows are no-ops under skip_duplicates. + Post-send failures (Read timeouts and errors) are ambiguous and are NOT + retried: the engine can abandon the transaction open on the pooled + connection, so a retry's statements stack into the same transaction and one + commit applies both increment sets. Ambiguous failures drop the batch; the + caller logs it at error. Callers must not add their own retry around this + function.""" if not transactions: return @@ -119,10 +129,14 @@ async def flush_tool_usage_transactions( key=lambda entry: (entry[0], entry[1]), ) + index_table: Final = SpendLogToolIndexRepository(prisma_client).table for attempt in range(n_retry_times + 1): try: + for statement_rows in spend_log_write_batches( + index_rows, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS + ): + await index_table.create_many(data=statement_rows, skip_duplicates=True) async with prisma_client.db.batch_() as batcher: - batcher.litellm_spendlogtoolindex.create_many(data=index_rows, skip_duplicates=True) for (date_key, tool_name), grouped in groupby(per_tool_day, key=lambda entry: (entry[0], entry[1])): entries = tuple(grouped) spend = sum(entry[2] for entry in entries) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index df967058cf0..797323794d2 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -17,8 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import billed_guardrail_cost_by_unit from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES +from litellm.proxy.db.spend_log_batching import spend_log_write_batches from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, @@ -401,13 +403,12 @@ async def process_spend_logs_guardrail_usage( return try: - # Insert index rows (skip duplicates by request_id + guardrail_id) - if index_rows: + index_table: Final = SpendLogGuardrailIndexRepository(prisma_client).table + for statement_rows in spend_log_write_batches( + index_rows, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS + ): try: - await SpendLogGuardrailIndexRepository(prisma_client).table.create_many( - data=index_rows, - skip_duplicates=True, - ) + await index_table.create_many(data=statement_rows, skip_duplicates=True) except Exception as e: verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 64da2ad00f6..3a61f773001 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -64,6 +64,7 @@ _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( CallTypes.pass_through.value, CallTypes.llm_passthrough_route.value, CallTypes.allm_passthrough_route.value, + CallTypes.call_mcp_tool.value, # CheckBatchCost's synthetic logging_obj for a completed managed batch carries # whatever LiteLLM_ManagedObjectTable stored at create time, and all of it is # None for a batch created before those columns were persisted, or by the master diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index cc594f167c7..e8aacac9e67 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -117,6 +117,14 @@ class ResponsesToolChatForm: web_search_options: OpenAIWebSearchOptions | None +@dataclass(frozen=True, slots=True) +class ResponsesReasoningChatForm: + """The Responses ``reasoning`` object as the two params Chat Completions takes.""" + + effort: str | None + summary: str | None + + if TYPE_CHECKING: from openai.types.responses.response_apply_patch_tool_call import ( ResponseApplyPatchToolCall, @@ -309,6 +317,79 @@ class LiteLLMCompletionResponsesConfig: ) return supported_params is not None and "web_search_options" not in supported_params + @staticmethod + def _completion_bridges_back_to_responses_api( + model: str, + custom_llm_provider: str | None, + tools: Sequence[ChatCompletionToolParam | OpenAIMcpServerTool] | None, + web_search_options: OpenAIWebSearchOptions | None, + reasoning_effort: str | None, + reasoning_summary: str | None, + api_base: str | None, + ) -> bool: + """ + Whether ``litellm.completion`` will route this model back onto the Responses API. + + Delegates to the same check ``litellm.completion`` itself runs, and is asked with the + params this transform is about to emit, so the two cannot reach different answers. + """ + from litellm.main import responses_api_bridge_check + + try: + model_info, _ = responses_api_bridge_check( + model=model, + custom_llm_provider=custom_llm_provider or "", + web_search_options=web_search_options, + tools=tools, + reasoning_effort=reasoning_effort, + reasoning_summary=reasoning_summary, + api_base=api_base, + ) + except Exception as e: # noqa: BLE001 # a capability probe must never fail the request it probes for + verbose_logger.debug("responses bridge: reasoning effort mode check failed: %s", e) + return False + return model_info.get("mode") == "responses" + + @staticmethod + def _transform_reasoning_for_chat_completion( + reasoning_param: Reasoning | str | None, + model: str, + custom_llm_provider: str | None, + tools: Sequence[ChatCompletionToolParam | OpenAIMcpServerTool] | None = None, + web_search_options: OpenAIWebSearchOptions | None = None, + api_base: str | None = None, + ) -> ResponsesReasoningChatForm: + """ + Split the Responses ``reasoning`` object into the params Chat Completions understands. + + ``reasoning_effort`` is a string enum there, so the object is never forwarded whole: a chat + provider either rejects it or silently drops it, and dropping it turns reasoning off while + still billing for the turn. ``summary`` has no chat equivalent, so it rides the + ``reasoning_summary`` alias, which ``litellm.completion`` reassembles into ``{effort, + summary}`` when it bridges the model back onto the Responses API, and is sent to nothing + else. + """ + if not reasoning_param: + return ResponsesReasoningChatForm(effort=None, summary=None) + if isinstance(reasoning_param, str): + return ResponsesReasoningChatForm(effort=reasoning_param, summary=None) + + effort: Final = reasoning_param.get("effort") + summary: Final = reasoning_param.get("summary") + if summary is None: + return ResponsesReasoningChatForm(effort=effort, summary=None) + + bridges_back: Final = LiteLLMCompletionResponsesConfig._completion_bridges_back_to_responses_api( + model=model, + custom_llm_provider=custom_llm_provider, + tools=tools, + web_search_options=web_search_options, + reasoning_effort=effort, + reasoning_summary=summary, + api_base=api_base, + ) + return ResponsesReasoningChatForm(effort=effort, summary=summary if bridges_back else None) + @staticmethod def transform_responses_api_request_to_chat_completion_request( model: str, @@ -339,23 +420,14 @@ class LiteLLMCompletionResponsesConfig: if text_param: response_format = LiteLLMCompletionResponsesConfig._transform_text_format_to_response_format(text_param) - # Extract reasoning_effort from reasoning parameter - reasoning_effort: Reasoning | str | None = None - reasoning_param: Final = responses_api_request.get("reasoning") - if reasoning_param: - if isinstance(reasoning_param, dict): - # reasoning can be {"effort": "low|medium|high", "summary": "detailed"} - # Keep the full dict when summary is set so the responses API bridge can - # forward it; otherwise use the effort string for chat completion (e.g. Gemini). - if "summary" in reasoning_param: - reasoning_effort = reasoning_param - elif "effort" in reasoning_param: - reasoning_effort = reasoning_param.get("effort") - else: - reasoning_effort = reasoning_param - elif isinstance(reasoning_param, str): - # reasoning could be a string directly - reasoning_effort = reasoning_param + reasoning: Final = LiteLLMCompletionResponsesConfig._transform_reasoning_for_chat_completion( + reasoning_param=responses_api_request.get("reasoning"), + model=model, + custom_llm_provider=custom_llm_provider, + tools=tools, + web_search_options=web_search_options, + api_base=kwargs.get("api_base"), + ) litellm_completion_request: dict = { "messages": LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( @@ -378,7 +450,8 @@ class LiteLLMCompletionResponsesConfig: "service_tier": kwargs.get("service_tier"), "web_search_options": web_search_options, "response_format": response_format, - "reasoning_effort": reasoning_effort, + "reasoning_effort": reasoning.effort, + "reasoning_summary": reasoning.summary, "context_management": responses_api_request.get("context_management"), # litellm specific params "custom_llm_provider": custom_llm_provider, diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 58f9dadc1ce..985d31af997 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -155,11 +155,9 @@ class MCPServer(BaseModel): access_groups: list[str] | None = None allow_all_keys: bool = False available_on_public_internet: bool = True - # Explicit opt-in to upstream-delegated authentication for ``oauth2`` - # servers. When ``auth_type == oauth2`` and this is ``True``, MCP requests - # bypass LiteLLM API-key/SSO auth (and the pre-emptive 401) so the client - # completes PKCE directly with the upstream MCP server. See - # ``MCPRequestHandler._target_servers_delegate_auth_to_upstream``. + # Legacy opt-in to upstream-delegated authentication for ``oauth2`` + # servers. LiteLLM admission still applies; use ``oauth_delegate`` for the + # supported client-forwarded OAuth flow. # # Honored only for ``auth_type == oauth2``; ignored for any other # ``auth_type``. OAuth pass-through for non-oauth2 servers diff --git a/litellm/utils.py b/litellm/utils.py index c6386471f28..7732cd88cb5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5563,12 +5563,21 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P split_model = strip_bedrock_routing_prefix(split_model) + provider_model_info: Final = ( + ProviderConfigManager.get_provider_model_info(model=split_model, provider=LlmProviders(custom_llm_provider)) + if custom_llm_provider in LlmProvidersSet + else None + ) + provider_cost_key: Final = ( + provider_model_info.get_model_cost_key(split_model) if provider_model_info is not None else None + ) + return PotentialModelNamesAndCustomLLMProvider( split_model=split_model, combined_model_name=combined_model_name, stripped_model_name=stripped_model_name, combined_stripped_model_name=combined_stripped_model_name, - provider_prefixed_model_name=provider_prefixed_model_name, + provider_prefixed_model_name=provider_cost_key or provider_prefixed_model_name, custom_llm_provider=cast(str, custom_llm_provider), ) diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index f853e9ff8d6..a7d4135d550 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -78,7 +78,7 @@ auth_family: none assertions: [succeeds] source: "mcp_server_manager.py:1485-1492" - rationale: Public/anonymous servers; delegate_auth_to_upstream + rationale: Explicitly anonymous true_passthrough servers - id: mcp.call_tool.none.succeeds module: mcp tier: P1 diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index eb4e893adb8..1553e788472 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -14,6 +14,18 @@ from litellm.constants import LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS from litellm.litellm_core_utils.logging_worker import LoggingWorker +class _RecordCollector(logging.Handler): + """Captures emitted log records so a test can assert on real logging output + (level, message args, traceback) instead of patching the logger object.""" + + def __init__(self) -> None: + super().__init__() + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record) + + class TestLoggingWorker: """Test cases for LoggingWorker functionality.""" @@ -26,9 +38,7 @@ class TestLoggingWorker: async def test_graceful_shutdown_with_clear_queue(self, logging_worker): """Test that cancellation triggers clear_queue to prevent 'never awaited' warnings.""" # Mock the clear_queue method to verify it's called during cancellation - with patch.object( - logging_worker, "clear_queue", new_callable=AsyncMock - ) as mock_clear_queue: + with patch.object(logging_worker, "clear_queue", new_callable=AsyncMock) as mock_clear_queue: # Start the worker logging_worker.start() @@ -195,9 +205,7 @@ class TestLoggingWorker: async def test_worker_handles_cancellation_gracefully(self, logging_worker): """Test that the worker handles cancellation without throwing exceptions.""" # Mock verbose_logger to capture debug messages - with patch( - "litellm.litellm_core_utils.logging_worker.verbose_logger" - ) as mock_logger: + with patch("litellm.litellm_core_utils.logging_worker.verbose_logger") as mock_logger: # Start the worker logging_worker.start() @@ -264,29 +272,21 @@ class TestLoggingWorker: small_worker._ensure_queue() # Mock verbose_logger to capture exception messages - with patch( - "litellm.litellm_core_utils.logging_worker.verbose_logger" - ) as mock_logger: + with patch("litellm.litellm_core_utils.logging_worker.verbose_logger") as mock_logger: # Fill the queue beyond capacity mock_coro = AsyncMock() for _ in range(5): # More than max_queue_size of 2 small_worker.enqueue(mock_coro()) # Should have logged queue full exceptions - exception_calls = [ - call - for call in mock_logger.exception.call_args_list - if "queue is full" in str(call) - ] + exception_calls = [call for call in mock_logger.exception.call_args_list if "queue is full" in str(call)] assert len(exception_calls) > 0 @pytest.mark.asyncio async def test_context_propagation(self, logging_worker): """Test that enqueued tasks execute in their original context.""" # Create a context variable for testing - test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar( - "test_context_var" - ) + test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar("test_context_var") # Track results from multiple tasks using asyncio.Event for synchronization task_results = [] @@ -364,36 +364,28 @@ class TestLoggingWorker: task_results.sort(key=lambda x: x["task_id"]) # Verify that each task saw its own context - assert ( - len(task_results) == 3 - ), f"Expected 3 results, got {len(task_results)}: {task_results}" + assert len(task_results) == 3, f"Expected 3 results, got {len(task_results)}: {task_results}" # Task 1 should see "context_1" task1_result = next((r for r in task_results if r["task_id"] == "task_1"), None) assert task1_result is not None, "Task 1 result not found" - assert ( - task1_result["context_accessible"] is True - ), "Task 1 should have access to context variable" - assert ( - task1_result["context_value"] == "context_1" - ), f"Task 1 should see 'context_1', got: {task1_result['context_value']}" + assert task1_result["context_accessible"] is True, "Task 1 should have access to context variable" + assert task1_result["context_value"] == "context_1", ( + f"Task 1 should see 'context_1', got: {task1_result['context_value']}" + ) # Task 2 should see "context_2" task2_result = next((r for r in task_results if r["task_id"] == "task_2"), None) assert task2_result is not None, "Task 2 result not found" - assert ( - task2_result["context_accessible"] is True - ), "Task 2 should have access to context variable" - assert ( - task2_result["context_value"] == "context_2" - ), f"Task 2 should see 'context_2', got: {task2_result['context_value']}" + assert task2_result["context_accessible"] is True, "Task 2 should have access to context variable" + assert task2_result["context_value"] == "context_2", ( + f"Task 2 should see 'context_2', got: {task2_result['context_value']}" + ) # Task 3 should not have access to the context variable task3_result = next((r for r in task_results if r["task_id"] == "task_3"), None) assert task3_result is not None, "Task 3 result not found" - assert ( - task3_result["context_accessible"] is False - ), "Task 3 should not have access to context variable" + assert task3_result["context_accessible"] is False, "Task 3 should not have access to context variable" @pytest.mark.asyncio async def test_semaphore_concurrency_limit(self): @@ -525,3 +517,182 @@ class TestLoggingWorker: asyncio.run(rebind_on_second_loop()) assert sorted(executed) == [0, 1, 2, 3, 4] + + @pytest.mark.asyncio + async def test_timeout_burst_logs_one_bounded_summary(self): + """Regression (LIT-7519): a burst of callback timeouts must log one bounded summary, + not a full ERROR traceback per timed-out callback. + + Before the fix every timed-out callback hit ``verbose_logger.exception`` in + ``_process_log_task``, so a slow Redis timing out many spend-tracking callbacks at once + produced one stacktrace each, clustered milliseconds apart. They must collapse into a + single WARNING that counts them, with no tracebacks. + """ + timeout_count = 25 + worker = LoggingWorker( + timeout=0.05, + max_queue_size=100, + concurrency=100, + timeout_summary_window=1.0, + ) + + async def slow_callback() -> None: + await asyncio.sleep(10) + + logger = logging.getLogger("LiteLLM") + collector = _RecordCollector() + previous_level = logger.level + logger.addHandler(collector) + logger.setLevel(logging.DEBUG) + try: + worker.start() + for _ in range(timeout_count): + worker.enqueue(slow_callback()) + + await worker.flush() + assert worker._timeout_summary_task is not None + await worker._timeout_summary_task + await worker.stop() + finally: + logger.removeHandler(collector) + logger.setLevel(previous_level) + + errors = [r for r in collector.records if r.levelno >= logging.ERROR] + warnings = [r for r in collector.records if r.levelno == logging.WARNING] + assert errors == [], "a timeout burst must not emit any ERROR tracebacks" + assert len(warnings) == 1, "the whole burst must collapse into one summary line" + summary = warnings[0].getMessage() + assert f"{timeout_count} callback(s) timed out" in summary + assert f"{timeout_count} timed out since start" in summary + assert "slow_callback" in summary + + @pytest.mark.asyncio + async def test_non_timeout_error_keeps_traceback(self): + """A real programming error in a callback must still log a full traceback, so the + timeout aggregation never hides genuine failures. + """ + worker = LoggingWorker(timeout=5.0, max_queue_size=10) + + async def failing_callback() -> None: + raise ValueError("boom") + + logger = logging.getLogger("LiteLLM") + collector = _RecordCollector() + previous_level = logger.level + logger.addHandler(collector) + logger.setLevel(logging.DEBUG) + try: + worker.start() + worker.enqueue(failing_callback()) + + await worker.flush() + await worker.stop() + finally: + logger.removeHandler(collector) + logger.setLevel(previous_level) + + errors = [r for r in collector.records if r.levelno >= logging.ERROR] + warnings = [r for r in collector.records if r.levelno == logging.WARNING] + assert len(errors) == 1, "a real error must still be logged once" + assert errors[0].exc_info is not None, "the traceback must be preserved" + assert warnings == [], "a single real error is not a timeout summary" + assert worker._timeout_summary_task is None + + @pytest.mark.asyncio + async def test_callback_raised_timeout_keeps_traceback(self): + """A callback that raises TimeoutError on its own did not hit the worker's deadline, + so it is a real failure and must keep its traceback rather than being folded into the + bounded burst summary. + """ + worker = LoggingWorker(timeout=5.0, max_queue_size=10, timeout_summary_window=1.0) + + async def raises_own_timeout() -> None: + raise asyncio.TimeoutError("callback's own downstream timeout") + + logger = logging.getLogger("LiteLLM") + collector = _RecordCollector() + previous_level = logger.level + logger.addHandler(collector) + logger.setLevel(logging.DEBUG) + try: + worker.start() + worker.enqueue(raises_own_timeout()) + + await worker.flush() + await worker.stop() + finally: + logger.removeHandler(collector) + logger.setLevel(previous_level) + + errors = [r for r in collector.records if r.levelno >= logging.ERROR] + warnings = [r for r in collector.records if r.levelno == logging.WARNING] + assert len(errors) == 1, "a callback-raised TimeoutError must still be logged once" + assert errors[0].exc_info is not None, "the traceback must be preserved" + assert warnings == [], "a callback-raised TimeoutError is not a worker-deadline timeout" + assert worker._timeout_summary_task is None + + @pytest.mark.asyncio + async def test_stop_flushes_pending_timeout_summary(self): + """A burst still inside its summary window when the worker stops must still emit its one + summary, instead of losing it when the event loop tears down. + """ + worker = LoggingWorker( + timeout=0.05, + max_queue_size=10, + concurrency=10, + timeout_summary_window=30.0, + ) + + async def slow_callback() -> None: + await asyncio.sleep(10) + + logger = logging.getLogger("LiteLLM") + collector = _RecordCollector() + previous_level = logger.level + logger.addHandler(collector) + logger.setLevel(logging.DEBUG) + try: + worker.start() + for _ in range(3): + worker.enqueue(slow_callback()) + + await worker.flush() + assert worker._timeout_summary_task is not None + assert not worker._timeout_summary_task.done(), "precondition: the summary window has not elapsed" + await worker.stop() + finally: + logger.removeHandler(collector) + logger.setLevel(previous_level) + + warnings = [r for r in collector.records if r.levelno == logging.WARNING] + assert len(warnings) == 1, "stop must flush the pending summary exactly once" + assert "3 callback(s) timed out" in warnings[0].getMessage() + assert worker._timeout_summary_task is None + + def test_loop_change_resets_timeout_summary_state(self): + """On an event-loop change the summary task is bound to the dead loop; it and the pending + burst count must reset so timeouts on the new loop arm a fresh summary rather than a stale, + stuck one that silently drops later summaries. + """ + worker = LoggingWorker(timeout=1.0, max_queue_size=10, timeout_summary_window=30.0) + + async def timed_out_callback() -> None: + await asyncio.sleep(10) + + async def arm_on_first_loop() -> None: + worker._ensure_queue() + coro = timed_out_callback() + worker._record_callback_timeout(coro) + coro.close() + assert worker._timeout_summary_task is not None + assert worker._timeout_burst_count == 1 + + asyncio.run(arm_on_first_loop()) + assert worker._timeout_summary_task is not None + + async def rebind_on_second_loop() -> None: + worker._ensure_queue() + assert worker._timeout_summary_task is None, "stale summary task must drop on loop change" + assert worker._timeout_burst_count == 0, "stale burst count must reset on loop change" + + asyncio.run(rebind_on_second_loop()) 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 8479397efa7..fb0311ef39b 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 @@ -1813,3 +1813,15 @@ def test_streaming_preserves_selected_model_for_private_accounting(): completion_response=assembled, custom_llm_provider="fireworks_ai", ) == pytest.approx(expected_cost) + + +@pytest.mark.parametrize( + "model, expected", + [ + ("deepseek-r1", "fireworks_ai/accounts/fireworks/models/deepseek-r1"), + ("glm-5p3-fast", "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast"), + ("accounts/fireworks/models/deepseek-r1", "fireworks_ai/accounts/fireworks/models/deepseek-r1"), + ], +) +def test_get_model_cost_key_resolves_short_names_to_long_keys(model: str, expected: str) -> None: + assert FireworksAIConfig().get_model_cost_key(model) == expected 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 c2f4f7163a0..90ce821d62e 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 @@ -1483,12 +1483,10 @@ class TestMCPOAuth2AuthFlow: as LiteLLM API keys, causing auth failures and empty tool listings. """ - async def test_oauth2_token_in_authorization_header_fallback(self): + async def test_oauth2_token_in_authorization_header_requires_litellm_admission(self): """ - When only the Authorization header is present with a non-LiteLLM OAuth2 - token AND the target server delegates auth to upstream, LiteLLM skips its - own validation entirely (so the upstream token is never mistaken for a - virtual key) and forwards the bearer upstream. + A bare Authorization token on the legacy delegated mode must establish + a LiteLLM principal rather than entering anonymously. """ from litellm.types.mcp import MCPAuth @@ -1510,6 +1508,7 @@ class TestMCPOAuth2AuthFlow: patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", new_callable=AsyncMock, + return_value=UserAPIKeyAuth(user_id="admitted-user"), ) as mock_auth, patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): @@ -1524,9 +1523,8 @@ class TestMCPOAuth2AuthFlow: ) = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) - # The upstream token is never validated as a LiteLLM key ... - mock_auth.assert_not_called() - # ... and is preserved for upstream forwarding. + assert auth_result.user_id == "admitted-user" + mock_auth.assert_awaited_once() assert oauth2_headers.get("Authorization") == "Bearer atlassian-oauth2-access-token-xyz" async def test_explicit_litellm_key_with_oauth2_authorization(self): @@ -2367,11 +2365,9 @@ class TestMCPDelegateAuthToUpstream: """ Tests for the ``delegate_auth_to_upstream`` per-server flag. - When set on an ``auth_type=oauth2`` MCP server, LiteLLM must skip its own - API-key/SSO check entirely so the client completes PKCE directly with the - upstream MCP server. The gate must fail closed for any non-oauth2 server, - any mixed-target request, and any request where the target cannot be - resolved. + The legacy flag no longer bypasses LiteLLM admission. OAuth discovery may + still use the anonymous cold-start challenge, but a presented bearer must + authenticate to LiteLLM unless a separate admission credential is supplied. """ @staticmethod @@ -2386,6 +2382,13 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=delegate_auth_to_upstream, ) + def test_legacy_delegate_cold_start_fails_closed_without_targets(self): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + _is_legacy_delegate_cold_start, + ) + + assert _is_legacy_delegate_cold_start(None, client_ip=None) is False + def test_build_mcp_server_table_preserves_delegate_auth_to_upstream(self): """Registry → API list rows must expose delegate_auth_to_upstream for the UI.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -2439,11 +2442,10 @@ class TestMCPDelegateAuthToUpstream: not_passthrough = passthrough.model_copy(update={"oauth_passthrough": False}) assert manager._build_mcp_server_table(not_passthrough).oauth_passthrough is False - async def test_delegate_skips_litellm_auth_with_no_authorization(self): + async def test_delegate_without_authorization_attempts_litellm_auth_before_cold_start(self): """ - oauth2 + delegate_auth_to_upstream=True, no Authorization header at - all → anonymous UserAPIKeyAuth and ``user_api_key_auth`` is never - called. + A credential-free discovery request attempts LiteLLM admission before + the route emits its RFC 9728 challenge. """ from litellm.types.mcp import MCPAuth @@ -2457,6 +2459,8 @@ class TestMCPDelegateAuthToUpstream: with ( patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + side_effect=HTTPException(status_code=401, detail="No key provided"), ) as mock_auth, patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): @@ -2466,17 +2470,12 @@ class TestMCPDelegateAuthToUpstream: ) auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() - async def test_delegate_with_upstream_token_in_authorization_skips_litellm_auth( - self, - ): + async def test_delegate_with_only_upstream_token_requires_litellm_auth(self): """ - oauth2 + delegate_auth_to_upstream=True with an upstream OAuth token in - ``Authorization``: the delegate gate fires before any LiteLLM validation, - so ``user_api_key_auth`` is never called and the bearer is forwarded - upstream untouched. Skipping the doomed validation is what keeps a tool - call that actually succeeds from carrying a phantom 401 auth span. + An upstream token cannot establish a LiteLLM principal and must not + reopen anonymous admission. """ from litellm.types.mcp import MCPAuth @@ -2491,6 +2490,7 @@ class TestMCPDelegateAuthToUpstream: patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", new_callable=AsyncMock, + side_effect=HTTPException(status_code=401, detail="Invalid API key"), ) as mock_auth, patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): @@ -2498,17 +2498,11 @@ class TestMCPDelegateAuthToUpstream: auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True, ) - ( - auth_result, - _, - _, - _, - oauth2_headers, - _, - ) = await MCPRequestHandler.process_mcp_request(scope) - assert isinstance(auth_result, UserAPIKeyAuth) - assert oauth2_headers.get("Authorization") == "Bearer upstream-pkce-token" - mock_auth.assert_not_called() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_awaited_once() async def test_delegate_off_still_requires_litellm_auth(self): """ @@ -2687,15 +2681,11 @@ class TestMCPDelegateAuthToUpstream: assert auth_result.user_id == "real-user" mock_auth.assert_called_once() - async def test_authorization_bearer_on_delegate_server_treated_as_upstream(self): + async def test_authorization_bearer_on_delegate_server_establishes_litellm_principal(self): """ - On a delegate server the ``Authorization`` header is, by contract, an - upstream token rather than a LiteLLM key — even when it is sk-shaped. It - is forwarded upstream without LiteLLM validation, so ``user_api_key_auth`` - is not called and no LiteLLM identity is resolved. Callers who need - LiteLLM identity / spend tracking on a delegate server must supply - ``x-litellm-api-key`` (see - test_explicit_litellm_key_takes_precedence_over_delegate). + A bare Authorization bearer now follows normal LiteLLM admission. A + separate x-litellm-api-key is required when Authorization is intended + for the upstream server. """ from litellm.types.mcp import MCPAuth @@ -2727,9 +2717,9 @@ class TestMCPDelegateAuthToUpstream: _, ) = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) - assert auth_result.user_id is None + assert auth_result.user_id == "real-user" assert oauth2_headers.get("Authorization") == "Bearer sk-1234" - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() async def test_delegate_ignored_for_client_credentials_server(self): """ @@ -2828,13 +2818,10 @@ class TestMCPDelegateAuthToUpstream: assert exc_info.value.status_code == 401 mock_auth.assert_called_once() - async def test_delegate_bypass_for_pure_pkce_server(self): + async def test_delegate_pkce_cold_start_attempts_litellm_auth(self): """ - oauth2 + delegate + oauth2_flow=None and NO stored client credentials - (pure PKCE, the common delegate case) → bypass must still fire. The - shape resolves to a non-M2M flow, so the security gate leaves it alone; - the fail-closed rule targets the M2M shape specifically, not every - unstamped row. + A pure PKCE server may defer a credential-free request to the route's + challenge, but normal LiteLLM admission still runs first. """ from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -2869,13 +2856,12 @@ class TestMCPDelegateAuthToUpstream: ): mock_mgr.get_mcp_server_by_name.return_value = pkce_server auth, *_rest = await MCPRequestHandler.process_mcp_request(scope) - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() assert auth.api_key is None - async def test_delegate_bypass_for_internal_server(self): + async def test_internal_delegate_cold_start_attempts_litellm_auth(self): """ - Delegate + oauth2 interactive servers bypass LiteLLM auth even when - ``available_on_public_internet`` is False (internal MCPs). + Internal delegated servers follow the same admission contract. """ from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -2910,14 +2896,11 @@ class TestMCPDelegateAuthToUpstream: ): mock_mgr.get_mcp_server_by_name.return_value = internal_server auth, *_rest = await MCPRequestHandler.process_mcp_request(scope) - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() assert auth.api_key is None - async def test_get_allowed_servers_excludes_client_credentials_delegate(self): - """ - get_allowed_mcp_servers must not surface M2M (client_credentials) delegate - servers to anonymous callers even if delegate_auth_to_upstream=True. - """ + async def test_get_allowed_servers_excludes_legacy_delegates(self): + """Legacy delegated servers are never added to anonymous access.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, ) @@ -2955,102 +2938,7 @@ class TestMCPDelegateAuthToUpstream: ): result = await manager.get_allowed_mcp_servers(None) - assert "pkce-server" in result - assert "m2m-server" not in result - - async def test_get_allowed_servers_excludes_unstamped_m2m_shape_delegate(self): - """ - The anonymous allow-list must also exclude an M2M-shape delegate server whose - oauth2_flow was never stamped (null column, verbatim-read as non-M2M). Reading - the bare has_client_credentials here would surface it to anonymous callers; the - resolved-flow check fails closed on the shape, matching the auth gate. - """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - manager = MCPServerManager() - pkce_server = MCPServer( - server_id="pkce-server", - name="pkce_server", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - available_on_public_internet=True, - ) - unstamped_m2m = MCPServer( - server_id="unstamped-m2m", - name="unstamped_m2m", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - oauth2_flow=None, - client_id="cid", - client_secret="csecret", - token_url="https://idp.example.com/token", - ) - assert unstamped_m2m.has_client_credentials is False - manager.registry = { - pkce_server.server_id: pkce_server, - unstamped_m2m.server_id: unstamped_m2m, - } - - with patch.object( - MCPRequestHandler, - "get_allowed_mcp_servers", - new_callable=AsyncMock, - return_value=[], - ): - result = await manager.get_allowed_mcp_servers(None) - - assert "pkce-server" in result - assert "unstamped-m2m" not in result - - async def test_get_allowed_servers_includes_internal_delegate(self): - """ - Internal-only (available_on_public_internet=False) delegate servers - appear in the anonymous allow-list like public delegate servers. - """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - manager = MCPServerManager() - public_server = MCPServer( - server_id="public-server", - name="public_server", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - available_on_public_internet=True, - ) - internal_server = MCPServer( - server_id="internal-server", - name="internal_server", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - available_on_public_internet=False, - ) - manager.registry = { - public_server.server_id: public_server, - internal_server.server_id: internal_server, - } - - with patch.object( - MCPRequestHandler, - "get_allowed_mcp_servers", - new_callable=AsyncMock, - return_value=[], - ): - result = await manager.get_allowed_mcp_servers(None) - - assert "public-server" in result - assert "internal-server" in result + assert result == [] async def test_true_passthrough_skips_litellm_auth_anonymously(self): """auth_type=true_passthrough performs no admission auth: the caller's Authorization is an diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 67948375403..28faf375ab8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -597,7 +597,11 @@ class TestHookHeaderMergePriority: "Authorization": "Bearer oauth2-token", "X-OAuth": "yes", }, - raw_headers=None, + raw_headers={ + "x-litellm-api-key": "Bearer sk-litellm-key", + "authorization": "Bearer oauth2-token", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), proxy_logging_obj=None, hook_extra_headers={ "Authorization": "Bearer hook-jwt", 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 7eb4396e60d..f5e4a420496 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 @@ -521,6 +521,93 @@ def test_prepare_mcp_server_headers_oauth2_interactive_drops_caller_authorizatio assert extra_headers is None +def test_prepare_mcp_server_headers_legacy_delegate_strips_admission_authorization(): + from litellm.proxy._experimental.mcp_server.server import ( + _prepare_mcp_server_headers, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="legacy-delegate-admission", + name="legacy-delegate", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=None, + mcp_auth_header=None, + oauth2_headers={"Authorization": "Bearer sk-litellm-key"}, + raw_headers={"authorization": "Bearer sk-litellm-key"}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), + ) + + assert server_auth_header is None + assert extra_headers is None + + +def test_prepare_mcp_server_headers_legacy_delegate_preserves_separate_upstream_authorization(): + from litellm.proxy._experimental.mcp_server.server import ( + _prepare_mcp_server_headers, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="legacy-delegate-dual-credential", + name="legacy-delegate", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=None, + mcp_auth_header=None, + oauth2_headers={"Authorization": "Bearer upstream-token"}, + raw_headers={ + "x-litellm-api-key": "Bearer sk-litellm-key", + "authorization": "Bearer upstream-token", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), + ) + + assert server_auth_header is None + assert extra_headers == {"Authorization": "Bearer upstream-token"} + + +def test_prepare_mcp_server_headers_legacy_delegate_strips_repeated_admission_key(): + from litellm.proxy._experimental.mcp_server.server import ( + _prepare_mcp_server_headers, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="legacy-delegate-repeated-key", + name="legacy-delegate", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=None, + mcp_auth_header=None, + oauth2_headers={"Authorization": "Bearer sk-litellm-key"}, + raw_headers={ + "x-litellm-api-key": "Bearer sk-litellm-key", + "authorization": "Bearer sk-litellm-key", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), + ) + + assert server_auth_header is None + assert extra_headers is None + + def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers(): """M2M must not merge caller Authorization from raw_headers when extra_headers lists it.""" try: @@ -6157,15 +6244,8 @@ def _patch_delegate_resolver(server: MCPServer, *resolvable_names: str): @pytest.mark.asyncio -async def test_delegate_bad_token_gets_connect_time_401(): - """Regression (LIT-4194): a rejected upstream token on a delegate-auth server - must fail the connect with 401 + ``error="invalid_token"``, not be absorbed - into HTTP 200 + an empty tool list by the tools/list handler. - - Delegate-mode clients send only ``Authorization`` (no ``x-litellm-api-key``), - so ``_get_forwarded_auth_from_scope`` returns None and, before the fix, the - preflight returned early without probing. - """ +async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-quality-ok: this removed security-sensitive egress has no return value; non-invocation is the contract + """A bare bearer is an admission credential and must never reach upstream.""" from litellm.proxy._experimental.mcp_server.server import ( _check_passthrough_upstream_auth, ) @@ -6184,48 +6264,6 @@ async def test_delegate_bad_token_gets_connect_time_401(): "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')), ) as probe, - ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( - scope=scope, - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, - ) - - assert exc_info.value.status_code == 401 - challenge = exc_info.value.headers["www-authenticate"] - assert 'error="invalid_token"' in challenge - assert ( - 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge - ) - probe.assert_awaited_once() - probe_url, probe_auth = probe.call_args.args - assert probe_url == "http://upstream:9401/mcp" - assert probe_auth == "Bearer bogus-token" - - -@pytest.mark.asyncio -async def test_delegate_valid_token_passes_preflight(): - """An upstream-accepted token must not be blocked by the delegate preflight.""" - from litellm.proxy._experimental.mcp_server.server import ( - _check_passthrough_upstream_auth, - ) - from litellm.proxy._types import UserAPIKeyAuth - - server = _delegate_auth_mcp_server() - scope = _delegate_scope([(b"authorization", b"Bearer good-token")]) - - with ( - _patch_delegate_resolver(server, "delegate_test"), - patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), - patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(200, None)), - ) as probe, ): await _check_passthrough_upstream_auth( scope=scope, @@ -6234,43 +6272,103 @@ async def test_delegate_valid_token_passes_preflight(): client_ip=None, ) - probe.assert_awaited_once() + probe.assert_not_awaited() @pytest.mark.asyncio -async def test_delegate_valid_token_forbidden_returns_403(): - """An upstream that accepts the token but forbids the caller (403) must surface - as a bare 403 with no ``WWW-Authenticate`` re-auth hint (a fresh token with the - same scopes would loop), not as an invalid_token challenge.""" +async def test_legacy_delegate_dual_credentials_are_not_probed_upstream(): # test-quality-ok: this removed security-sensitive egress has no return value; non-invocation is the contract + """A separate upstream bearer never triggers the removed legacy probe.""" from litellm.proxy._experimental.mcp_server.server import ( _check_passthrough_upstream_auth, ) from litellm.proxy._types import UserAPIKeyAuth server = _delegate_auth_mcp_server() - scope = _delegate_scope([(b"authorization", b"Bearer scoped-out-token")]) + scope = _delegate_scope( + [ + (b"x-litellm-api-key", b"sk-litellm-proxy-key"), + (b"authorization", b"Bearer upstream-token"), + ] + ) with ( - _patch_delegate_resolver(server, "delegate_test"), - patch( + patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), - patch( + patch( # test-quality-ok: the removed probe call is the security regression under test "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(403, None)), - ), + new=AsyncMock(), + ) as probe, ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.parametrize( + "probe_status, expected_status", + [(200, None), (401, 401), (403, 403)], +) +@pytest.mark.asyncio +async def test_oauth_passthrough_preflight_preserves_status_contract(probe_status, expected_status): + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="passthrough-id", + name="passthrough_server", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + oauth_passthrough=True, + extra_headers=["Authorization"], + ) + scope = _delegate_scope( + [ + (b"x-litellm-api-key", b"sk-litellm-proxy-key"), + (b"authorization", b"Bearer upstream-token"), + ] + ) + + with ( + patch( # test-quality-ok: isolate authorized-server resolution so this test exercises the preflight contract + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( # test-quality-ok: the upstream transport boundary is the behavior being mapped to an HTTP response + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(probe_status, None)), + ) as probe, + ): + if expected_status is None: + result = await _check_passthrough_upstream_auth( scope=scope, - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], + user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), + mcp_servers=["passthrough_server"], client_ip=None, ) + assert result is None + else: + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), + mcp_servers=["passthrough_server"], + client_ip=None, + ) + assert exc_info.value.status_code == expected_status + if expected_status == 401: + assert "passthrough_server" in exc_info.value.headers["www-authenticate"] - assert exc_info.value.status_code == 403 - assert not (exc_info.value.headers or {}) + probe.assert_awaited_once_with("https://upstream.example.com/mcp", "Bearer upstream-token") @pytest.mark.asyncio @@ -6428,123 +6526,6 @@ async def test_delegate_not_probed_when_named_only_via_server_id(): probe.assert_not_awaited() -@pytest.mark.asyncio -async def test_delegate_preflight_with_unpatched_probe(): - """Integration across the preflight and the unpatched ``_probe_upstream_auth``, - mocked only at the httpx-client boundary (tests/test_litellm is mocked-only; the - real-network proof lives in the PR's live-proxy evidence). The mock honors the - ``AsyncHTTPHandler.post`` contract by raising ``httpx.HTTPStatusError`` on the - upstream 401, so the production ``except httpx.HTTPStatusError`` branch is the one - exercised. A rejected token surfaces as the connect-time 401 challenge; an - accepted token passes untouched, and the caller's bearer reaches the delegate URL.""" - import httpx - - from litellm.proxy._experimental.mcp_server.server import ( - _check_passthrough_upstream_auth, - ) - from litellm.proxy._types import UserAPIKeyAuth - - accepted = MagicMock() - accepted.status_code = 200 - accepted.headers = {} - rejected = MagicMock() - rejected.status_code = 401 - rejected.headers = {"www-authenticate": 'Bearer realm="stub-upstream", error="invalid_token"'} - - async def respond_by_token(url=None, headers=None, json=None, timeout=None, **kwargs): - if headers.get("Authorization") == "Bearer good-token": - return accepted - raise httpx.HTTPStatusError( - "401 Unauthorized", - request=httpx.Request("POST", url), - response=rejected, - ) - - mock_client = MagicMock() - mock_client.post = AsyncMock(side_effect=respond_by_token) - - server = _delegate_auth_mcp_server() - - with ( - _patch_delegate_resolver(server, "delegate_test"), - patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), - patch( - "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", - return_value=mock_client, - ), - ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( - scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]), - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, - ) - - await _check_passthrough_upstream_auth( - scope=_delegate_scope([(b"authorization", b"Bearer good-token")]), - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, - ) - - assert exc_info.value.status_code == 401 - challenge = exc_info.value.headers["www-authenticate"] - assert 'error="invalid_token"' in challenge - assert ( - 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge - ) - probed_urls = [call.kwargs["url"] for call in mock_client.post.await_args_list] - assert probed_urls == ["http://upstream:9401/mcp", "http://upstream:9401/mcp"] - - -@pytest.mark.asyncio -async def test_delegate_challenge_echoes_requested_alias(): - """An alias-routed delegate request must be probed, and the challenge must echo - the requested alias (not the canonical server name) so the resource_metadata - URL matches what the tokenless preemptive challenge emits for the same route.""" - from litellm.proxy._experimental.mcp_server.server import ( - _check_passthrough_upstream_auth, - ) - from litellm.proxy._types import UserAPIKeyAuth - - server = _delegate_auth_mcp_server().model_copy(update={"alias": "dt-alias"}) - scope = { - "type": "http", - "method": "POST", - "path": "/mcp/dt-alias", - "scheme": "http", - "server": ("localhost", 4000), - "headers": [(b"authorization", b"Bearer bogus-token")], - } - - with ( - _patch_delegate_resolver(server, "dt-alias"), - patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), - patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')), - ), - ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( - scope=scope, - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["dt-alias"], - client_ip=None, - ) - - challenge = exc_info.value.headers["www-authenticate"] - assert 'error="invalid_token"' in challenge - assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/dt-alias"' in challenge - - @pytest.mark.asyncio async def test_delegate_probe_not_fanned_out_to_access_group_members(): """A single access-group name passes the one-target route gate but must not fan @@ -6578,41 +6559,6 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members(): probe.assert_not_awaited() -def test_is_delegate_upstream_probe_target_fails_closed_on_m2m_shape(): - """An unstamped M2M-shape row (null ``oauth2_flow`` + client credentials) - resolves to ``client_credentials`` and must not be probed with the caller's - bearer; its stored client credentials drive egress instead.""" - from litellm.proxy._experimental.mcp_server.server import ( - _is_delegate_upstream_probe_target, - ) - - assert _is_delegate_upstream_probe_target(_delegate_auth_mcp_server()) is True - - m2m_shape = MCPServer( - server_id="delegate-m2m", - name="delegate_m2m", - url="http://upstream:9401/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - oauth2_flow=None, - token_url="http://idp:9000/token", - client_id="client", - client_secret="secret", - ) - assert _is_delegate_upstream_probe_target(m2m_shape) is False - - non_delegate = MCPServer( - server_id="oauth2-plain", - name="oauth2_plain", - url="http://upstream:9401/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - oauth2_flow="authorization_code", - ) - assert _is_delegate_upstream_probe_target(non_delegate) is False - - @pytest.mark.asyncio async def test_create_mcp_client_sampling_disabled_by_default(): """Sampling callback must be None when allow_sampling is not set (default False).""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d2987c5112e..d56f08c4e79 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -3390,6 +3390,72 @@ class TestMCPServerManager: ) assert not extra_headers or "authorization" not in {k.lower() for k in extra_headers} + @pytest.mark.asyncio + async def test_call_regular_mcp_tool_legacy_delegate_never_forwards_admission_key(self): + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="server-legacy-delegate-leak", + name="legacy-delegate", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + extra_headers = await self._capture_call_extra_headers( + server, + oauth2_headers={"Authorization": "Bearer sk-litellm-key"}, + raw_headers={"authorization": "Bearer sk-litellm-key"}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), + ) + assert not extra_headers or "authorization" not in {k.lower() for k in extra_headers} + + @pytest.mark.asyncio + async def test_call_regular_mcp_tool_legacy_delegate_forwards_separate_authorization(self): + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="server-legacy-delegate-dual-credential", + name="legacy-delegate", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + extra_headers = await self._capture_call_extra_headers( + server, + oauth2_headers={"Authorization": "Bearer upstream-token"}, + raw_headers={ + "x-litellm-api-key": "Bearer sk-litellm-key", + "authorization": "Bearer upstream-token", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), + ) + assert extra_headers == {"Authorization": "Bearer upstream-token"} + + @pytest.mark.asyncio + async def test_call_regular_mcp_tool_legacy_delegate_strips_repeated_admission_key(self): + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="server-legacy-delegate-repeated-key", + name="legacy-delegate", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + extra_headers = await self._capture_call_extra_headers( + server, + oauth2_headers={"Authorization": "Bearer sk-litellm-key"}, + raw_headers={ + "x-litellm-api-key": "Bearer sk-litellm-key", + "authorization": "Bearer sk-litellm-key", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), + ) + assert not extra_headers or "authorization" not in {k.lower() for k in extra_headers} + def test_should_strip_caller_authorization_new_modes(self): from litellm.proxy._types import UserAPIKeyAuth @@ -6916,51 +6982,6 @@ class TestMCPServerManager: == expected_server_ids ) - @pytest.mark.asyncio - async def test_get_allowed_mcp_servers_anonymous_delegate_requires_oauth2(self): - """Anonymous delegated auth listing should only include oauth2 servers.""" - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, - ) - - manager = MCPServerManager() - oauth_delegate_server = MCPServer( - server_id="oauth-delegate", - name="oauth_delegate", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - ) - api_key_delegate_server = MCPServer( - server_id="api-key-delegate", - name="api_key_delegate", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, - delegate_auth_to_upstream=True, - ) - oauth_non_delegate_server = MCPServer( - server_id="oauth-non-delegate", - name="oauth_non_delegate", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=False, - ) - manager.registry = { - oauth_delegate_server.server_id: oauth_delegate_server, - api_key_delegate_server.server_id: api_key_delegate_server, - oauth_non_delegate_server.server_id: oauth_non_delegate_server, - } - - with patch.object( - MCPRequestHandler, - "get_allowed_mcp_servers", - new_callable=AsyncMock, - return_value=[], - ): - result = await manager.get_allowed_mcp_servers(None) - - assert set(result) == {"oauth-delegate"} - def test_get_mcp_server_from_tool_name_uses_server_name_not_name(self): """ Test that _get_mcp_server_from_tool_name uses server.server_name instead of server.name @@ -8088,9 +8109,9 @@ class TestMCPServerTokenExchangeColumns: assert rebuilt_table.token_exchange_profile == "entra_obo" -class TestInternalDelegatePkceWarningLog: +class TestLegacyDelegateAuthWarningLog: @pytest.mark.asyncio - async def test_build_mcp_server_logs_on_internal_delegate_interactive(self, caplog): + async def test_build_mcp_server_logs_deprecation_for_internal_delegate(self, caplog): caplog.set_level(logging.WARNING, logger="LiteLLM") manager = MCPServerManager() table_record = LiteLLM_MCPServerTable( @@ -8106,11 +8127,11 @@ class TestInternalDelegatePkceWarningLog: ) await manager.build_mcp_server_from_table(table_record) combined = " ".join(r.getMessage() for r in caplog.records) - assert "internal-only" in combined + assert "deprecated auth_type=oauth2" in combined assert "delegate_auth_to_upstream=true" in combined @pytest.mark.asyncio - async def test_build_mcp_server_no_internal_delegate_log_when_public(self, caplog): + async def test_build_mcp_server_logs_deprecation_for_public_delegate(self, caplog): caplog.set_level(logging.WARNING, logger="LiteLLM") manager = MCPServerManager() table_record = LiteLLM_MCPServerTable( @@ -8126,12 +8147,13 @@ class TestInternalDelegatePkceWarningLog: ) await manager.build_mcp_server_from_table(table_record) combined = " ".join(r.getMessage() for r in caplog.records) - assert "internal-only" not in combined + assert "deprecated auth_type=oauth2" in combined + assert "auth_type=oauth_delegate" in combined def test_warn_skipped_for_client_credentials(self, caplog): caplog.set_level(logging.WARNING, logger="LiteLLM") from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _warn_internal_delegate_pkce_if_applicable, + _warn_legacy_delegate_auth_if_applicable, ) server = MCPServer( @@ -8144,9 +8166,9 @@ class TestInternalDelegatePkceWarningLog: available_on_public_internet=False, delegate_auth_to_upstream=True, ) - _warn_internal_delegate_pkce_if_applicable(server, source="test") + _warn_legacy_delegate_auth_if_applicable(server, source="test") combined = " ".join(r.getMessage() for r in caplog.records) - assert "internal-only" not in combined + assert "deprecated auth_type=oauth2" not in combined class TestHasClientCredentialsOAuth2Flow: diff --git a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py index 9c4fbbf41aa..282c2a7cfaa 100644 --- a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py +++ b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py @@ -1,15 +1,17 @@ """ Tests for the tool usage writer: ToolUsageTransaction construction (invoked tools -only) and the flush that writes LiteLLM_SpendLogToolIndex plus the -LiteLLM_DailyToolSpend rollup in one transaction. +only) and the flush that writes LiteLLM_SpendLogToolIndex in bounded statements +plus the LiteLLM_DailyToolSpend rollup in one transaction. """ from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +from litellm.proxy.db import spend_log_tool_index from litellm.proxy.db.spend_log_tool_index import ( ToolUsageTransaction, build_tool_usage_transaction, @@ -35,11 +37,24 @@ class _FakeBatcher: return None +def _prisma(batch_: MagicMock) -> MagicMock: + prisma = MagicMock() + prisma.db.batch_ = batch_ + prisma.db.litellm_spendlogtoolindex.create_many = AsyncMock() + return prisma + + def _prisma_with_batcher() -> tuple[MagicMock, _FakeBatcher]: batcher = _FakeBatcher() - prisma = MagicMock() - prisma.db.batch_ = MagicMock(return_value=batcher) - return prisma, batcher + return _prisma(MagicMock(return_value=batcher)), batcher + + +def _index_rows_written(prisma: MagicMock) -> list[tuple[str, str]]: + return [ + (row["request_id"], row["tool_name"]) + for call in prisma.db.litellm_spendlogtoolindex.create_many.call_args_list + for row in call.kwargs["data"] + ] class TestBuildToolUsageTransaction: @@ -228,9 +243,8 @@ class TestFlushToolUsageTransactions: prisma_client=prisma, transactions=[_transaction("r1", tool_names=("tool_a", "tool_b"), spend=0.10, total_tokens=100)], ) - index_rows = batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["data"] - assert [(r["request_id"], r["tool_name"]) for r in index_rows] == [("r1", "tool_a"), ("r1", "tool_b")] - assert batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["skip_duplicates"] is True + assert _index_rows_written(prisma) == [("r1", "tool_a"), ("r1", "tool_b")] + assert prisma.db.litellm_spendlogtoolindex.create_many.call_args.kwargs["skip_duplicates"] is True upserts = { c.kwargs["where"]["date_tool_name"]["tool_name"]: c.kwargs["data"] @@ -267,19 +281,46 @@ class TestFlushToolUsageTransactions: assert data["update"]["request_count"] == {"increment": 2} @pytest.mark.asyncio - async def test_index_rows_and_rollup_share_one_transaction(self): - # Both writes go through the same batch_() so a failed flush cannot leave - # index rows without their rollup increments (or vice versa); increments - # are not idempotent, so partial states must be unreachable. + async def test_index_rows_are_written_in_bounded_statements_outside_the_rollup_transaction(self, monkeypatch): + monkeypatch.setattr(spend_log_tool_index, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100) prisma, batcher = _prisma_with_batcher() - await flush_tool_usage_transactions( - prisma_client=prisma, - transactions=[_transaction("r1")], - ) + tool_names = tuple(f"tool_{i}" for i in range(50)) + transactions = [_transaction(f"r{i}", tool_names=tool_names) for i in range(5)] + await flush_tool_usage_transactions(prisma_client=prisma, transactions=transactions) + + statements = prisma.db.litellm_spendlogtoolindex.create_many.call_args_list + assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50] + assert all(call.kwargs["skip_duplicates"] is True for call in statements) + assert _index_rows_written(prisma) == [ + (txn.request_id, tool_name) for txn in transactions for tool_name in tool_names + ] + batcher.litellm_spendlogtoolindex.create_many.assert_not_called() + prisma.db.batch_.assert_called_once() + assert batcher.litellm_dailytoolspend.upsert.call_count == len(tool_names) + + @pytest.mark.asyncio + async def test_index_connection_error_is_retried_before_the_rollup_is_attempted(self, monkeypatch): + prisma, batcher = _prisma_with_batcher() + prisma.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=[httpx.ConnectError("down"), None]) + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep) + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + assert prisma.db.litellm_spendlogtoolindex.create_many.await_count == 2 prisma.db.batch_.assert_called_once() - batcher.litellm_spendlogtoolindex.create_many.assert_called_once() batcher.litellm_dailytoolspend.upsert.assert_called_once() + @pytest.mark.asyncio + async def test_ambiguous_index_error_drops_the_batch_without_touching_the_rollup(self): + prisma, _ = _prisma_with_batcher() + prisma.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=httpx.ReadTimeout("ambiguous")) + with pytest.raises(httpx.ReadTimeout): + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + prisma.db.litellm_spendlogtoolindex.create_many.assert_awaited_once() + prisma.db.batch_.assert_not_called() + @pytest.mark.asyncio async def test_empty_batch_touches_nothing(self): prisma, _ = _prisma_with_batcher() @@ -290,11 +331,8 @@ class TestFlushToolUsageTransactions: async def test_connection_errors_retry_and_succeed(self, monkeypatch): # A failed batch commits nothing, so retrying a connection error cannot # double-count; the flush must retry rather than drop the batch. - import httpx - batcher = _FakeBatcher() - prisma = MagicMock() - prisma.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), batcher]) + prisma = _prisma(MagicMock(side_effect=[httpx.ConnectError("down"), batcher])) sleeps: list[float] = [] async def fake_sleep(seconds: float) -> None: @@ -308,10 +346,7 @@ class TestFlushToolUsageTransactions: @pytest.mark.asyncio async def test_connection_errors_exhaust_retries_then_raise(self, monkeypatch): - import httpx - - prisma = MagicMock() - prisma.db.batch_ = MagicMock(side_effect=httpx.ConnectError("down")) + prisma = _prisma(MagicMock(side_effect=httpx.ConnectError("down"))) async def fake_sleep(seconds: float) -> None: return None @@ -325,8 +360,7 @@ class TestFlushToolUsageTransactions: @pytest.mark.asyncio async def test_non_connection_errors_do_not_retry(self): - prisma = MagicMock() - prisma.db.batch_ = MagicMock(side_effect=ValueError("bad data")) + prisma = _prisma(MagicMock(side_effect=ValueError("bad data"))) with pytest.raises(ValueError, match="bad data"): await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) prisma.db.batch_.assert_called_once() @@ -338,11 +372,8 @@ class TestFlushToolUsageTransactions: # unknown; the engine can leave the transaction open on the pooled # connection, so a retry's statements would stack into it and one # commit would apply both increment sets. These must never retry. - import httpx - error = getattr(httpx, ambiguous_error)("ambiguous") - prisma = MagicMock() - prisma.db.batch_ = MagicMock(side_effect=error) + prisma = _prisma(MagicMock(side_effect=error)) with pytest.raises((httpx.ReadTimeout, httpx.ReadError)): await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) prisma.db.batch_.assert_called_once() diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 110de7dbe70..85f22e1f307 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest +from litellm.proxy.guardrails import usage_tracking from litellm.proxy.guardrails.usage_tracking import ( _MAX_PENDING_ROWS, PendingRollups, @@ -511,3 +512,56 @@ async def test_requeued_cost_is_added_to_the_next_flush(): costs = _cost_upserts(recovered) assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0) assert costs["someFutureCounter"] == (0.0, 7) + + +def _fan_out_payload(request_id: str, guardrail_ids: tuple[str, ...]) -> dict[str, Any]: + return { + "request_id": request_id, + "startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc), + "team_id": "team-a", + "api_key": "hashed-key-1", + "metadata": json.dumps( + {"guardrail_information": [{"guardrail_id": gid, "guardrail_status": "success"} for gid in guardrail_ids]} + ), + } + + +def _index_rows_written(prisma: MagicMock) -> list[tuple[str, str]]: + return [ + (row["request_id"], row["guardrail_id"]) + for call in prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list + for row in call.kwargs["data"] + ] + + +@pytest.mark.asyncio +async def test_index_rows_are_written_in_row_bounded_statements(monkeypatch): + """ + LIT-5931: the drain caps logs, not logs x guardrails, so a fan-out must be + split into statements the query engine can afford instead of one create_many. + """ + monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100) + prisma = _prisma() + guardrail_ids = tuple(f"guard-{i}" for i in range(50)) + logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)] + + await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups()) + + statements = prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list + assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50] + assert all(call.kwargs["skip_duplicates"] is True for call in statements) + assert _index_rows_written(prisma) == [(f"r{i}", gid) for i in range(5) for gid in guardrail_ids] + + +@pytest.mark.asyncio +async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollup(monkeypatch): + monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100) + prisma = _prisma() + prisma.db.litellm_spendlogguardrailindex.create_many.side_effect = [None, httpx.ReadTimeout("ambiguous"), None] + guardrail_ids = tuple(f"guard-{i}" for i in range(50)) + logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)] + + await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups()) + + assert prisma.db.litellm_spendlogguardrailindex.create_many.await_count == 3 + assert prisma.db.litellm_dailyguardrailmetrics.upsert.await_count == len(guardrail_ids) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index fad137af66b..395ce68ec54 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -19,6 +19,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( run_spend_event, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.utils import ProxyUpdateSpend from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -1912,14 +1913,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit(): ("allm_passthrough_route", True), ("aretrieve_batch", True), ("acompletion", False), - ("call_mcp_tool", False), + ("call_mcp_tool", True), (None, False), ], ) def test_should_track_cost_callback_pass_through_without_owner(call_type, expected): """Regression for LIT-3782: unauthenticated pass-through requests (auth=false) carry no key/user/team/end-user, yet must still be tracked so they land in - LiteLLM_SpendLogs. Other call types with no owner stay untracked. + LiteLLM_SpendLogs. Explicit MCP passthrough calls require the same handling. + Other call types with no owner stay untracked. aretrieve_batch is included for the same reason: CheckBatchCost's synthetic logging_obj for a completed managed batch only ever carries @@ -1939,10 +1941,26 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ) +def test_should_track_cost_callback_respects_disabled_spend_updates(monkeypatch): + monkeypatch.setattr(ProxyUpdateSpend, "disable_spend_updates", staticmethod(lambda: True)) + + assert ( + _should_track_cost_callback( + user_api_key="key", + user_id="user", + team_id="team", + end_user_id="end-user", + call_type="call_mcp_tool", + ) + is False + ) + + @pytest.mark.parametrize( "call_type, expect_spend_log", [ ("pass_through_endpoint", True), + ("call_mcp_tool", True), ("aretrieve_batch", True), ("acompletion", False), (None, False), @@ -1953,8 +1971,8 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(cal """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It - must now be written for pass-through call types while other unauthenticated - calls remain skipped. + must now be written for pass-through and MCP tool call types while other + unauthenticated calls remain skipped. aretrieve_batch is included because CheckBatchCost's completed-batch cost event reaches this same callback with no attributable key/user/team when diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 342ec4435a7..3c78bbf79d7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2527,6 +2527,180 @@ class TestToolTransformation: "type": "object", } + @pytest.mark.parametrize( + "model, custom_llm_provider", + [ + ("bedrock/converse/global.anthropic.claude-sonnet-5", "bedrock_converse"), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"), + ("claude-sonnet-5", "vertex_ai"), + ("gemini-3.1-pro-preview", "vertex_ai"), + ("moonshotai.kimi-k2-thinking", "bedrock_mantle"), + ], + ) + def test_reasoning_summary_still_yields_a_string_reasoning_effort(self, model, custom_llm_provider): + """ + A Responses request carrying ``reasoning.summary`` must still reach a chat provider as a + plain ``reasoning_effort`` string. ``summary`` is Responses-only, and forwarding the whole + object turns reasoning off: Bedrock Converse and Vertex silently discard a non-string + ``reasoning_effort``, and Bedrock Mantle rejects the request outright. + """ + responses_api_request = {"reasoning": {"effort": "medium", "summary": "auto"}} + + result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input="hi", + responses_api_request=responses_api_request, + custom_llm_provider=custom_llm_provider, + ) + + assert result["reasoning_effort"] == "medium" + + @pytest.mark.parametrize( + "model, custom_llm_provider", + [ + ("gpt-5.4-pro", "azure_ai"), + ("gpt-5", "openai"), + ("gpt-5.1", "openai"), + ("gpt-5", "azure"), + ], + ) + def test_bridged_model_carries_the_summary_as_an_alias(self, model, custom_llm_provider): + """ + ``summary`` reaches a bridged model through the ``reasoning_summary`` alias, never smuggled + inside ``reasoning_effort``. ``litellm.completion`` reads that alias back with + ``peek_reasoning_summary_aliases`` and reassembles ``{effort, summary}``, so the far end + gets the same object it always did while no chat provider ever sees a non-string effort. + """ + result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input="hi", + responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}}, + custom_llm_provider=custom_llm_provider, + ) + + assert result["reasoning_effort"] == "medium" + assert result["reasoning_summary"] == "auto" + + @pytest.mark.parametrize( + "model, custom_llm_provider", + [ + ("gpt-5", "openai"), + ("gpt-5.1", "openai"), + ("gpt-5", "azure"), + ], + ) + def test_gpt_5_summary_survives_the_bridge_it_claims_to_take(self, model, custom_llm_provider): + """ + Regression for the probe disagreeing with the real decision. The transform asked + ``responses_api_bridge_check`` with ``reasoning_summary`` taken straight off the Responses + object, but ``litellm.completion`` reads it from ``optional_params`` via + ``peek_reasoning_summary_aliases``, which the bridged request never populated. So these + models answered "bridging" to the probe and "not bridging" for real, and the object landed + on Chat Completions, which only takes a string. Emitting the alias makes the two agree. + """ + from litellm.main import responses_api_bridge_check + from litellm.utils import get_optional_params, peek_reasoning_summary_aliases + + result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input="hi", + responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}}, + custom_llm_provider=custom_llm_provider, + ) + optional_params = get_optional_params( + model=model, + custom_llm_provider=custom_llm_provider, + reasoning_effort=result["reasoning_effort"], + reasoning_summary=result["reasoning_summary"], + ) + model_info, _ = responses_api_bridge_check( + model=model, + custom_llm_provider=custom_llm_provider, + reasoning_effort=result["reasoning_effort"], + reasoning_summary=peek_reasoning_summary_aliases(optional_params), + ) + + assert model_info.get("mode") == "responses" + + def test_a_failing_bridge_probe_falls_back_to_the_string_effort(self, monkeypatch): + """ + The probe is a capability question, so a model-info lookup blowing up must not fail the + request. It degrades to the chat-safe form: a string effort and no alias. + """ + import litellm.main + + def _boom(**_kwargs): + raise RuntimeError("model info unavailable") + + monkeypatch.setattr(litellm.main, "responses_api_bridge_check", _boom) + + result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gpt-5.4-pro", + input="hi", + responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}}, + custom_llm_provider="azure_ai", + ) + + assert result["reasoning_effort"] == "medium" + assert "reasoning_summary" not in result + + @pytest.mark.parametrize( + "reasoning, expected", + [ + ({"effort": "high"}, "high"), + ("low", "low"), + ({"summary": "auto"}, None), + ({}, None), + (None, None), + ], + ) + def test_reasoning_param_shapes_map_to_reasoning_effort(self, reasoning, expected): + """ + An object without ``effort`` carries nothing Chat Completions can use, so no + ``reasoning_effort`` is sent at all (the bridge drops None-valued params). + """ + result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + input="hi", + responses_api_request={"reasoning": reasoning}, + custom_llm_provider="bedrock", + ) + + assert result.get("reasoning_effort") == expected + assert ("reasoning_effort" in result) is (expected is not None) + + @pytest.mark.parametrize( + "model, expected_thinking", + [ + ("global.anthropic.claude-sonnet-5", {"type": "adaptive"}), + ( + "anthropic.claude-sonnet-4-5-20250929-v1:0", + {"type": "enabled", "budget_tokens": 2048}, + ), + ], + ) + def test_reasoning_summary_still_enables_thinking_on_bedrock(self, model, expected_thinking): + """ + End to end through Bedrock Converse's own param mapping: the effort a Responses request asks + for must survive into ``thinking``, whether the model takes an adaptive effort or a legacy + token budget. Forwarding the object instead leaves ``thinking`` unset and the model never + reasons, which is the failure this guards. + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + bridged = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input="hi", + responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}}, + custom_llm_provider="bedrock", + ) + + mapped = AmazonConverseConfig().map_openai_params( + {"reasoning_effort": bridged["reasoning_effort"]}, {}, model, True + ) + + assert expected_thinking.items() <= mapped["thinking"].items() + def test_bedrock_anthropic_responses_tools_yield_only_function_toolspec(self): """ End-to-end (no network) of the LIT-3858 acceptance criterion: the mixed tools array diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index fc4207fb26b..d89496a0fd0 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4449,6 +4449,75 @@ def test_fireworks_models_in_backup_cost_map(): ), f"short-form {short_key} does not match long-form {long_key}" +@pytest.fixture +def fireworks_short_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr( + litellm, + "model_cost", + { + "fireworks_ai/accounts/fireworks/models/glm-5p3": { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "litellm_provider": "fireworks_ai", + "mode": "chat", + "max_tokens": 100, + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "input_cost_per_token": 2.1e-6, + "output_cost_per_token": 6.6e-6, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "fireworks_ai/nomic-ai/nomic-embed-text-v1.5": { + "input_cost_per_token": 8e-9, + "output_cost_per_token": 0.0, + "litellm_provider": "fireworks_ai", + "mode": "embedding", + }, + }, + ) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def test_fireworks_short_model_names_resolve_to_long_cost_map_keys(fireworks_short_model_cost_map: None) -> None: + model_info = litellm.get_model_info("fireworks_ai/glm-5p3") + assert model_info["key"] == "fireworks_ai/accounts/fireworks/models/glm-5p3" + assert model_info["input_cost_per_token"] == 1e-6 + assert model_info["max_tokens"] == 100 + + model_info = litellm.get_model_info("glm-5p3", custom_llm_provider="fireworks_ai") + assert model_info["key"] == "fireworks_ai/accounts/fireworks/models/glm-5p3" + + model_info = litellm.get_model_info("fireworks_ai/glm-5p3-fast") + assert model_info["key"] == "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast" + assert model_info["input_cost_per_token"] == 2.1e-6 + + model_info = litellm.get_model_info("fireworks_ai/nomic-ai/nomic-embed-text-v1.5") + assert model_info["key"] == "fireworks_ai/nomic-ai/nomic-embed-text-v1.5" + + with pytest.raises(Exception, match="isn't mapped"): + litellm.get_model_info("fireworks_ai/does-not-exist") + + +def test_fireworks_short_model_names_price_with_completion_cost(fireworks_short_model_cost_map: None) -> None: + from litellm.types.utils import ModelResponse + + response = ModelResponse( + model="fireworks_ai/glm-5p3", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + cost = litellm.completion_cost( + completion_response=response, + model="fireworks_ai/glm-5p3", + custom_llm_provider="fireworks_ai", + ) + + assert cost == pytest.approx(10 * 1e-6 + 5 * 2e-6) + + class TestBedrockBaseModelLabelKeepsTools: """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly label must not silently drop ``tools``/``tool_choice`` under ``drop_params``."""