scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_fix_dotprompt_model_swap

Browse source
This commit is contained in:
mateo-berri 2026-08-26 17:36:03 -07:00
commit 8697a9ffa9
28 changed files with 1938 additions and 299 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

83
litellm/cost_calculator.py
View file

@ -2,6 +2,7 @@
## File for 'response_cost' calculation in Logging
import logging
import time
from collections.abc import Sequence
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Final, Literal, cast
@ -2373,6 +2374,64 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
def _candidate_realtime_token_costs(
model_name: str,
combined_usage_object: Usage,
custom_llm_provider: str,
data_residency: str | None,
) -> tuple[float, float] | None:
try:
return generic_cost_per_token(
model=model_name,
usage=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
except Exception:
return None
def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) -> bool:
entries: Final = (
litellm.model_cost.get(model_name),
litellm.model_cost.get(f"{custom_llm_provider}/{model_name}"),
)
return any(
entry is not None and any("cost_per" in field and value is not None for field, value in entry.items())
for entry in entries
)
def _first_priced_realtime_token_costs(
potential_model_names: Sequence[str | None],
combined_usage_object: Usage,
custom_llm_provider: str,
data_residency: str | None,
) -> tuple[float, float]:
candidate_costs: Final = (
(model_name, costs)
for model_name in potential_model_names
if model_name is not None
and (
costs := _candidate_realtime_token_costs(
model_name=model_name,
combined_usage_object=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
)
is not None
)
return next(
(
costs
for model_name, costs in candidate_costs
if sum(costs) > 0 or _cost_map_entry_declares_pricing(model_name, custom_llm_provider)
),
(0.0, 0.0),
)
def handle_realtime_stream_cost_calculation(
results: OpenAIRealtimeStreamList,
combined_usage_object: Usage,
@ -2397,24 +2456,12 @@ def handle_realtime_stream_cost_calculation(
potential_model_names.append(received_model)
potential_model_names.append(litellm_model_name)
input_cost_per_token = 0.0
output_cost_per_token = 0.0
for model_name in potential_model_names:
try:
if model_name is None:
continue
_input_cost_per_token, _output_cost_per_token = generic_cost_per_token(
model=model_name,
usage=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
except Exception:
continue
input_cost_per_token += _input_cost_per_token
output_cost_per_token += _output_cost_per_token
break # exit if we find a valid model
input_cost_per_token, output_cost_per_token = _first_priced_realtime_token_costs(
potential_model_names=potential_model_names,
combined_usage_object=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
transcription_cost: Final = (
handle_realtime_transcription_cost_calculation(
results=results,

147
litellm/model_prices_and_context_window_backup.json
View file

@ -20150,7 +20150,7 @@
"gemini-3.5-flash-lite": {
"deprecation_date": "2027-07-21",
"cache_read_input_token_cost": 3e-08,
"cache_read_input_token_cost_flex": 2e-08,
"cache_read_input_token_cost_flex": 1.5e-08,
"cache_read_input_token_cost_priority": 5e-08,
"input_cost_per_token": 3e-07,
"input_cost_per_token_batches": 1.5e-07,
@ -20332,7 +20332,7 @@
"supports_image_size": false
},
"gemini-2.5-flash-preview-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost": 3e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-language-models",
@ -20377,10 +20377,53 @@
"google_maps_grounding_cost_per_query": 0.025,
"supports_image_size": false
},
"gemini-live-2.5-flash-native-audio": {
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "realtime",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supported_endpoints": [
"/vertex_ai/live"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.035,
"search_context_size_medium": 0.035,
"search_context_size_high": 0.035
},
"gemini_native_audio": true
},
"gemini-live-2.5-flash-preview-native-audio-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
@ -20424,7 +20467,7 @@
"gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
@ -20469,7 +20512,7 @@
},
"gemini-2.5-flash-lite-preview-06-17": {
"deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost": 1e-08,
"input_cost_per_audio_token": 5e-07,
"input_cost_per_token": 1e-07,
"litellm_provider": "vertex_ai-language-models",
@ -21126,18 +21169,15 @@
},
"gemini-2.5-pro-preview-tts": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_above_200k_tokens": 2.5e-07,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"input_cost_per_token": 1e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_above_200k_tokens": 1.5e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview",
"output_cost_per_token": 2e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_modalities": [
"text"
],
@ -22062,7 +22102,7 @@
"supports_image_size": false
},
"gemini/gemini-2.5-flash-preview-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost": 3e-08,
"deprecation_date": "2026-02-17",
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
@ -22111,7 +22151,7 @@
"supports_image_size": false
},
"gemini/gemini-flash-latest": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost": 3e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"litellm_provider": "gemini",
@ -22158,7 +22198,7 @@
"google_maps_grounding_cost_per_query": 0.025
},
"gemini/gemini-flash-lite-latest": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost": 1e-08,
"input_cost_per_audio_token": 3e-07,
"input_cost_per_token": 1e-07,
"litellm_provider": "gemini",
@ -22206,7 +22246,7 @@
},
"gemini/gemini-2.5-flash-lite-preview-06-17": {
"deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost": 1e-08,
"input_cost_per_audio_token": 5e-07,
"input_cost_per_token": 1e-07,
"litellm_provider": "gemini",
@ -22254,11 +22294,11 @@
"supports_image_size": false
},
"gemini/gemini-2.5-flash-preview-tts": {
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"mode": "audio_speech",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_token": 1e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/audio/speech"
],
@ -23212,19 +23252,16 @@
},
"gemini/gemini-2.5-pro-preview-tts": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_above_200k_tokens": 2.5e-07,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"input_cost_per_token": 1e-06,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_above_200k_tokens": 1.5e-05,
"output_cost_per_token": 2e-05,
"rpm": 10000,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview",
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_modalities": [
"text"
],
@ -42027,7 +42064,7 @@
"vertex_ai/gemini-3.5-flash-lite": {
"deprecation_date": "2027-07-21",
"cache_read_input_token_cost": 3e-08,
"cache_read_input_token_cost_flex": 2e-08,
"cache_read_input_token_cost_flex": 1.5e-08,
"cache_read_input_token_cost_priority": 5e-08,
"input_cost_per_token": 3e-07,
"input_cost_per_token_batches": 1.5e-07,
@ -48884,15 +48921,16 @@
}
},
"gemini-2.5-flash-native-audio-latest": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -48909,15 +48947,16 @@
"gemini_native_audio": true
},
"gemini-2.5-flash-native-audio-preview-09-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -48934,15 +48973,16 @@
"gemini_native_audio": true
},
"gemini-2.5-flash-native-audio-preview-12-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -48992,15 +49032,16 @@
"gemini_audio_only_live": true
},
"gemini/gemini-2.5-flash-native-audio-latest": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -49019,15 +49060,16 @@
"gemini_native_audio": true
},
"gemini/gemini-2.5-flash-native-audio-preview-09-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -49046,15 +49088,16 @@
"gemini_native_audio": true
},
"gemini/gemini-2.5-flash-native-audio-preview-12-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -49123,11 +49166,11 @@
"rpm": 10
},
"gemini-2.5-flash-preview-tts": {
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"mode": "audio_speech",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_token": 1e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/audio/speech"
]

4
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -1600,7 +1600,7 @@ async def refresh_user_oauth_token(
) -> OAuthCredentialPayload | None:
"""Attempt to refresh a per-user OAuth2 token using its stored refresh_token.
POSTs to ``server.token_url`` with ``grant_type=refresh_token``.
POSTs to ``server.effective_token_url`` with ``grant_type=refresh_token``.
On success: persists the new credential via ``store_user_oauth_credential``
and returns the updated payload dict.
@ -1609,7 +1609,7 @@ async def refresh_user_oauth_token(
stale credential and triggering re-authentication.
"""
refresh_token: Final[str | None] = cred.get("refresh_token")
token_url: Final[str | None] = getattr(server, "token_url", None)
token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None)
server_id: Final[str] = getattr(server, "server_id", "")
client_id: Final[str | None] = getattr(server, "client_id", None)
client_secret: Final[str | None] = getattr(server, "client_secret", None)

183
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -3,7 +3,7 @@ import html as _html
import json
import secrets
import time
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
@ -663,6 +663,26 @@ def _endpoint_not_configured_detail(
)
async def _server_with_oauth_endpoints(
mcp_server: MCPServer,
needed_endpoint: Callable[[MCPServer], str | None],
) -> MCPServer:
"""Join deferred OAuth discovery only when the endpoint this caller needs is still missing.
Admin-entered endpoints live on ``configured_*`` after an anchored issuer empties the
resolved fields. A caller whose needed endpoint already resolves never awaits discovery
and cannot 503 over a leftover pin. A server still missing it joins the deferred task;
no slot is a no-op and the caller 400s.
"""
if needed_endpoint(mcp_server) is not None:
return mcp_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load
global_mcp_server_manager,
)
return await global_mcp_server_manager.ensure_oauth_metadata_discovered(mcp_server)
def _raise_unless_oauth2_discovery_server(
mcp_server: MCPServer | None,
mcp_server_name: str | None,
@ -697,7 +717,7 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool:
returns directly to the client's redirect URI without transiting the gateway. Gateway-side
redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit
arm, where the upstream only knows the gateway's own callback."""
return mcp_server.is_dcr_bridge and bool(mcp_server.registration_url) and not mcp_server.client_id
return mcp_server.is_dcr_bridge and bool(mcp_server.effective_registration_url) and not mcp_server.client_id
def _require_s256_pkce(
@ -745,7 +765,7 @@ def _redirect_to_upstream_authorize(
**({"scope": scope_value} if scope_value else {}),
**({"resource": upstream_resource} if upstream_resource else {}),
}
parsed_auth_url: Final = urlparse(mcp_server.authorization_url or "")
parsed_auth_url: Final = urlparse(mcp_server.effective_authorization_url or "")
merged_params: Final = {**dict(parse_qsl(parsed_auth_url.query)), **passthrough_params}
return RedirectResponse(urlunparse(parsed_auth_url._replace(query=urlencode(merged_params))))
@ -812,18 +832,19 @@ async def authorize_with_server(
ephemeral_dcr_client: "EphemeralDcrClient | None" = None,
):
_raise_if_not_oauth2(mcp_server)
if mcp_server.authorization_url is None:
resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint)
if resolved_server.effective_authorization_url is None:
raise HTTPException(
status_code=400,
detail=_endpoint_not_configured_detail(
mcp_server,
resolved_server,
"authorization url",
"set Authorization URL and Token URL manually",
"set Issuer to discover them from the identity provider (RFC 8414)",
),
)
if mcp_server.is_dcr_bridge:
if resolved_server.is_dcr_bridge:
# Enforce S256 PKCE on both bridge arms. The relay arm forwards the validated,
# now-non-optional pair to the upstream authorize; the short-circuit arm keeps
# calling this for its enforcement side effect, then falls through to the gateway
@ -832,9 +853,9 @@ async def authorize_with_server(
# A gateway-minted ephemeral client is registered against {base}/callback, so its
# flow must run the short-circuit arm; the relay arm is only for clients that
# registered themselves through the front door and hold their own redirect binding.
if _dcr_bridge_relays_client_registration(mcp_server) and ephemeral_dcr_client is None:
if _dcr_bridge_relays_client_registration(resolved_server) and ephemeral_dcr_client is None:
return _redirect_to_upstream_authorize(
mcp_server=mcp_server,
mcp_server=resolved_server,
client_id=client_id,
redirect_uri=redirect_uri,
state=state,
@ -860,7 +881,7 @@ async def authorize_with_server(
# litellm key, so the browser session is the only identity source; without one there is nothing to
# bind, so send the user through login first. Every other oauth2 server keeps the identity-less state.
litellm_user_id: str | None = None
if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate:
if resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate:
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import
_user_id_from_session_cookie,
)
@ -870,7 +891,7 @@ async def authorize_with_server(
return _redirect_to_litellm_login(request)
denial: Final = await _bridge_authorize_access_denial(
litellm_user_id=litellm_user_id,
mcp_server=mcp_server,
mcp_server=resolved_server,
redirect_uri=redirect_uri,
state=state,
)
@ -884,7 +905,7 @@ async def authorize_with_server(
code_challenge_method=code_challenge_method,
client_redirect_uri=redirect_uri,
litellm_user_id=litellm_user_id,
mcp_server_id=mcp_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None,
mcp_server_id=resolved_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None,
dcr_client_id=ephemeral_dcr_client.client_id if ephemeral_dcr_client else None,
dcr_client_secret=ephemeral_dcr_client.client_secret if ephemeral_dcr_client else None,
dcr_token_endpoint_auth_method=ephemeral_dcr_client.token_endpoint_auth_method
@ -894,26 +915,26 @@ async def authorize_with_server(
relay_state: Final = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES)
params: Final = {
"client_id": mcp_server.client_id if mcp_server.client_id else client_id,
"client_id": resolved_server.client_id if resolved_server.client_id else client_id,
"redirect_uri": f"{request_base_url}/callback",
"state": relay_state,
"response_type": response_type or "code",
}
if scope:
params["scope"] = scope
elif mcp_server.scopes:
params["scope"] = " ".join(mcp_server.scopes)
elif resolved_server.scopes:
params["scope"] = " ".join(resolved_server.scopes)
if code_challenge:
params["code_challenge"] = code_challenge
if code_challenge_method:
params["code_challenge_method"] = code_challenge_method
upstream_resource: Final = resolve_upstream_resource(mcp_server)
upstream_resource: Final = resolve_upstream_resource(resolved_server)
if upstream_resource:
params["resource"] = upstream_resource
parsed_auth_url: Final = urlparse(mcp_server.authorization_url)
parsed_auth_url: Final = urlparse(resolved_server.effective_authorization_url)
existing_params: Final = dict(parse_qsl(parsed_auth_url.query))
existing_params.update(params)
final_url: Final = urlunparse(parsed_auth_url._replace(query=urlencode(existing_params)))
@ -946,11 +967,13 @@ async def exchange_token_with_server(
if grant_type not in ("authorization_code", "refresh_token"):
raise HTTPException(status_code=400, detail="Unsupported grant_type")
if mcp_server.token_url is None:
resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _token_flow_needed_endpoint)
token_url: Final = resolved_server.effective_token_url
if token_url is None:
raise HTTPException(
status_code=400,
detail=_endpoint_not_configured_detail(
mcp_server,
resolved_server,
"token url",
"set Token URL manually",
"set Issuer to discover it from the identity provider (RFC 8414)",
@ -965,16 +988,16 @@ async def exchange_token_with_server(
# recovered from a sealed code) must authenticate the way its own registration was granted,
# not the way the server row is configured; callers that carry no method keep the row's method
# as before.
resolved_client_id: Final = mcp_server.client_id if mcp_server.client_id else client_id
resolved_client_secret: Final = mcp_server.client_secret if mcp_server.client_id else client_secret
resolved_client_id: Final = resolved_server.client_id if resolved_server.client_id else client_id
resolved_client_secret: Final = resolved_server.client_secret if resolved_server.client_id else client_secret
resolved_auth_method: Final = (
mcp_server.token_endpoint_auth_method
if mcp_server.client_id
else (client_token_endpoint_auth_method or mcp_server.token_endpoint_auth_method)
resolved_server.token_endpoint_auth_method
if resolved_server.client_id
else (client_token_endpoint_auth_method or resolved_server.token_endpoint_auth_method)
)
try:
token_request: Final = build_upstream_oauth2_token_request(
mcp_server,
resolved_server,
auth_method=resolved_auth_method,
client_id=resolved_client_id,
client_secret=resolved_client_secret,
@ -987,14 +1010,14 @@ async def exchange_token_with_server(
bridge_upstream_refresh: SecretStr | None = None
bridge_upstream_scope: str | None = None
refresh_request_scope: str | None = None
is_bridge: Final = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge
is_bridge: Final = resolved_server.is_oauth_delegate and resolved_server.is_dcr_bridge
if grant_type == "refresh_token":
# Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed
# identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange
# sends the upstream token and never the envelope. A failure returns without touching the upstream.
if is_bridge:
prepared_refresh: Final = await _prepare_bridge_refresh(mcp_server, refresh_token)
prepared_refresh: Final = await _prepare_bridge_refresh(resolved_server, refresh_token)
if not isinstance(prepared_refresh, _BridgeRefreshReady):
return _bridge_mint_error_response(prepared_refresh)
bridge_mint_ready = prepared_refresh.ready
@ -1031,13 +1054,13 @@ async def exchange_token_with_server(
# A raw upstream code (scripted path) opens to None and the code is used as-is.
bridge_identity = open_bridge_authorization_code(code)
if bridge_identity is not None:
if bridge_identity.mcp_server_id != mcp_server.server_id:
if bridge_identity.mcp_server_id != resolved_server.server_id:
raise HTTPException(
status_code=400,
detail="Authorization code was issued for a different MCP server",
)
code = bridge_identity.upstream_code
bridge_token_relay: Final = _dcr_bridge_relays_client_registration(mcp_server)
bridge_token_relay: Final = _dcr_bridge_relays_client_registration(resolved_server)
if bridge_token_relay and not redirect_uri:
raise HTTPException(
status_code=400,
@ -1059,7 +1082,7 @@ async def exchange_token_with_server(
# Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or
# the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code.
if is_bridge:
prepared: Final = await _prepare_bridge_mint(request, mcp_server, bridge_identity)
prepared: Final = await _prepare_bridge_mint(request, resolved_server, bridge_identity)
if not isinstance(prepared, _BridgeMintReady):
return _bridge_mint_error_response(prepared)
bridge_mint_ready = prepared
@ -1067,7 +1090,7 @@ async def exchange_token_with_server(
async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
try:
response: Final = await async_client.post(
mcp_server.token_url,
token_url,
headers={"Accept": "application/json", **token_request.headers},
data=token_data,
)
@ -1076,8 +1099,8 @@ async def exchange_token_with_server(
except httpx.HTTPStatusError as exc:
fault: Final = classify_upstream_token_rejection(
exc.response,
credential_source=_token_credential_source(mcp_server),
log_context=mcp_server.server_id,
credential_source=_token_credential_source(resolved_server),
log_context=resolved_server.server_id,
)
upstream_rejected_bridge_refresh: Final = (
is_bridge
@ -1090,7 +1113,7 @@ async def exchange_token_with_server(
"bridge refresh: the upstream rejected the sealed refresh token for server=%s with "
"invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client "
"re-runs authorization_code rather than an opaque upstream error",
mcp_server.server_id,
resolved_server.server_id,
)
return _bridge_mint_error_response("invalid_refresh")
return render_token_fault(fault)
@ -1103,22 +1126,22 @@ async def exchange_token_with_server(
# Validate token response against server-configured rules before any storage.
# This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc.
if mcp_server.token_validation and isinstance(mcp_server.token_validation, dict):
if resolved_server.token_validation and isinstance(resolved_server.token_validation, dict):
_validate_token_response(
token_response=token_response,
validation_rules=mcp_server.token_validation,
server_id=mcp_server.server_id,
validation_rules=resolved_server.token_validation,
server_id=resolved_server.server_id,
)
# Store server-side when the server is configured for per-user OAuth and
# the calling client has provided a valid LiteLLM identity.
# Errors are non-fatal: the token is still returned to the client.
if mcp_server.needs_user_oauth_token:
if resolved_server.needs_user_oauth_token:
user_id: Final = await _extract_user_id_from_request(request)
if user_id:
try:
await _store_per_user_token_server_side(
server=mcp_server,
server=resolved_server,
user_id=user_id,
token_response=token_response,
)
@ -1126,7 +1149,7 @@ async def exchange_token_with_server(
verbose_logger.warning(
"exchange_token_with_server: server-side storage failed for user=%s server=%s: %s",
user_id,
mcp_server.server_id,
resolved_server.server_id,
exc,
)
else:
@ -1136,7 +1159,7 @@ async def exchange_token_with_server(
"requires the stored token, so the client will be challenged with 401 on reconnect. "
"Ensure the request carries a valid LiteLLM key (x-litellm-api-key or Authorization), "
"or store it via POST /mcp/server/{id}/oauth-user-credential.",
mcp_server.server_id,
resolved_server.server_id,
)
# A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the
@ -1147,7 +1170,9 @@ async def exchange_token_with_server(
token_response = {**token_response, "scope": refresh_request_scope}
# Phase 3: seal the upstream grant into the client-held envelope; failures map through the same
# OAuth-shaped response as the phase-1 preconditions.
minted: Final = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc))
minted: Final = _finish_bridge_mint(
bridge_mint_ready, resolved_server, token_response, datetime.now(timezone.utc)
)
return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted)
raw_access_token: Final = token_response.get("access_token") if isinstance(token_response, dict) else None
@ -1551,7 +1576,8 @@ async def mint_ephemeral_dcr_client(request: Request, mcp_server: MCPServer) ->
bounded by the server count even when the request origin varies) so parallel authorize requests
cannot each register an upstream client; the cache stamps nothing onto the server record and
correctness never depends on it because the sealed state carries the client through the flow."""
if mcp_server.registration_url is None:
registration_url: Final = mcp_server.effective_registration_url
if registration_url is None:
return None
request_base_url: Final = get_request_base_url(request)
cache_key: Final = f"mcp_ephemeral_dcr_client:{mcp_server.server_id}:{request_base_url}"
@ -1571,7 +1597,7 @@ async def mint_ephemeral_dcr_client(request: Request, mcp_server: MCPServer) ->
"token_endpoint_auth_method": "none",
}
response: Final = await _post_dcr_registration(
registration_url=mcp_server.registration_url,
registration_url=registration_url,
register_data=register_data,
server_id=mcp_server.server_id,
)
@ -1617,7 +1643,7 @@ async def resolve_ephemeral_dcr_client(
usable to generate orphan IdP clients)."""
if not (mcp_server.is_true_passthrough or (mcp_server.is_oauth_delegate and not mcp_server.is_dcr_bridge)):
return None
if mcp_server.authorization_url is None:
if mcp_server.effective_authorization_url is None:
raise HTTPException(
status_code=400,
detail="MCP server authorization url is not set",
@ -1627,6 +1653,29 @@ async def resolve_ephemeral_dcr_client(
return await mint_ephemeral_dcr_client(request, mcp_server)
def _register_flow_needed_endpoint(mcp_server: MCPServer) -> str | None:
"""The register flow's deferred-discovery join gate. A DCR bridge with no admin-configured
client can only register callers through the upstream's registration endpoint
(``_oauth_endpoints_unresolved`` keeps its discovery slot armed for exactly this shape), so
the flow must keep joining discovery while registration is still missing instead of silently
degrading to the dummy short-circuit. Every other shape only needs the authorization url."""
if mcp_server.is_dcr_bridge and not mcp_server.client_id and mcp_server.effective_registration_url is None:
return None
return mcp_server.effective_authorization_url
def _token_flow_needed_endpoint(mcp_server: MCPServer) -> str | None:
"""The token exchange's deferred-discovery join gate. The exchange's relay-vs-callback arm
(:func:`_dcr_bridge_relays_client_registration`) reads the registration url, so a clientless
DCR bridge rebuilt without its discovered registration endpoint must keep joining discovery
even when the token url already resolves; skipping it would select the gateway-callback arm
and the upstream would reject the code over a redirect_uri mismatch. Every other shape only
needs the token url."""
if mcp_server.is_dcr_bridge and not mcp_server.client_id and mcp_server.effective_registration_url is None:
return None
return mcp_server.effective_token_url
async def register_client_with_server(
request: Request,
mcp_server: MCPServer,
@ -1661,21 +1710,23 @@ async def register_client_with_server(
):
return dummy_return
if mcp_server.authorization_url is None:
resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint)
if resolved_server.effective_authorization_url is None:
raise HTTPException(
status_code=400,
detail=_endpoint_not_configured_detail(
mcp_server,
resolved_server,
"authorization url",
"set Authorization URL and Token URL manually",
"set Issuer to discover them from the identity provider (RFC 8414)",
),
)
if mcp_server.registration_url is None:
registration_url: Final = resolved_server.effective_registration_url
if registration_url is None:
return dummy_return
bridge_relay: Final = _dcr_bridge_relays_client_registration(mcp_server)
bridge_relay: Final = _dcr_bridge_relays_client_registration(resolved_server)
if bridge_relay and not client_redirect_uris:
raise HTTPException(
status_code=400,
@ -1690,15 +1741,17 @@ async def register_client_with_server(
"token_endpoint_auth_method": token_endpoint_auth_method or ("none" if bridge_relay else ""),
}
response: Final = await _post_dcr_registration(
registration_url=mcp_server.registration_url,
registration_url=registration_url,
register_data=register_data,
server_id=mcp_server.server_id,
server_id=resolved_server.server_id,
)
token_response = response.json()
if persist_credentials and not bridge_relay:
persistence_result = await _persist_dcr_client_registration(mcp_server, token_response, current_redirect_uri)
persistence_result = await _persist_dcr_client_registration(
resolved_server, token_response, current_redirect_uri
)
if persistence_result == "reused":
return dummy_return
@ -1755,17 +1808,10 @@ async def authorize(
lookup_name: Final[str | None] = mcp_server_name or client_id
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = (
await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip)
if lookup_name
else None
global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None
)
if mcp_server is None and mcp_server_name is None:
unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
mcp_server = (
await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server)
if unresolved_server is not None
else None
)
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
_raise_if_not_oauth2(mcp_server)
@ -1846,14 +1892,9 @@ async def token_endpoint(
lookup_name: Final = mcp_server_name or client_id
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip)
if mcp_server is None and mcp_server_name is None:
unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
mcp_server = (
await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server)
if unresolved_server is not None
else None
)
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
return await exchange_token_with_server(
@ -2684,10 +2725,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
return await register_aggregate_client(request=request, request_body=data)
resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
resolved_server: Final = await global_mcp_server_manager.ensure_oauth_metadata_discovered(resolved)
return await register_client_with_server(
request=request,
mcp_server=resolved_server,
mcp_server=resolved,
client_name=data.get("client_name", ""),
grant_types=data.get("grant_types", []),
response_types=data.get("response_types", []),
@ -2697,10 +2737,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
)
return dummy_return
mcp_server: Final = await global_mcp_server_manager.get_resolved_mcp_server_by_name(
mcp_server_name,
client_ip=client_ip,
)
mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip)
if mcp_server is None:
return dummy_return
return await register_client_with_server(

14
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -523,7 +523,7 @@ def _oauth_endpoints_unresolved(server: MCPServer) -> bool:
# can come from resource discovery, so a server that resolved its endpoints but no scopes is
# still unresolved for its flow.
return True
if server.is_dcr_bridge and not server.client_id and server.registration_url is None:
if server.is_dcr_bridge and not server.client_id and server.effective_registration_url is None:
# A DCR bridge with no admin-configured client can only register callers through the
# upstream's registration endpoint, so a build that resolved the authorize and token
# endpoints but not registration_endpoint (partial metadata) is still unresolved for its
@ -535,8 +535,8 @@ def _oauth_endpoints_unresolved(server: MCPServer) -> bool:
return _flow_endpoints_missing(
server.auth_type,
MCPServerManager.effective_oauth2_flow(server),
server.authorization_url,
server.token_url,
server.effective_authorization_url,
server.effective_token_url,
server.token_exchange_endpoint,
)
@ -6205,14 +6205,6 @@ class MCPServerManager:
return server
return None
async def get_resolved_mcp_server_by_name(
self,
server_name: str,
client_ip: str | None = None,
) -> MCPServer | None:
server: Final = self.get_mcp_server_by_name(server_name, client_ip=client_ip)
return await self.ensure_oauth_metadata_discovered(server) if server is not None else None
def get_filtered_registry(self, client_ip: str | None = None) -> dict[str, MCPServer]:
"""
Get registry filtered by client IP access control.

13
litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py
View file

@ -67,7 +67,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
rest of the identity rather than stored in a key."""
material: Final = "\x00".join(
(
server.token_url or "",
server.effective_token_url or "",
server.client_id or "",
server.client_secret or "",
" ".join(server.scopes or ()),
@ -82,7 +82,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
@staticmethod
def _has_client_credentials_config(server: "MCPServer") -> bool:
return bool(server.client_id and server.client_secret and server.token_url)
return bool(server.client_id and server.client_secret and server.effective_token_url)
async def async_get_token(self, server: "MCPServer") -> str | None:
"""Return a valid access token, fetching or refreshing as needed.
@ -112,19 +112,20 @@ class MCPOAuth2TokenCache(InMemoryCache):
return token
async def _fetch_token(self, server: "MCPServer") -> tuple[str, int]:
"""POST to ``token_url`` with ``grant_type=client_credentials``.
"""POST to ``effective_token_url`` with ``grant_type=client_credentials``.
Returns ``(access_token, ttl_seconds)`` where ttl accounts for the
expiry buffer so the cache entry expires before the real token does.
"""
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
if not server.client_id or not server.client_secret or not server.token_url:
token_url: Final = server.effective_token_url
if not server.client_id or not server.client_secret or not token_url:
raise ValueError(
f"MCP server '{server.server_id}' missing required OAuth2 fields: "
f"client_id={bool(server.client_id)}, "
f"client_secret={bool(server.client_secret)}, "
f"token_url={bool(server.token_url)}"
f"token_url={bool(token_url)}"
)
token_request: Final = build_upstream_oauth2_token_request(
@ -146,7 +147,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
)
try:
response: Final = await client.post(server.token_url, data=data, headers=token_request.headers or None)
response: Final = await client.post(token_url, data=data, headers=token_request.headers or None)
response.raise_for_status()
except httpx.HTTPStatusError as exc:
raise ValueError(

4
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
View file

@ -142,7 +142,7 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
config=ClientCredentialsConfig(
client_id=server.client_id,
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
token_url=server.token_url,
token_url=server.effective_token_url,
scopes=tuple(server.scopes or ()),
audience=server.audience,
upstream_resource=resolve_upstream_resource(server),
@ -163,7 +163,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
forwarded only when the operator set it; a missing one is omitted, not derived.
"""
endpoint: Final = server.token_exchange_endpoint or server.token_url
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
if not server.client_id or not server.client_secret:
return None
profile: Final[Literal["rfc8693", "entra_obo"]] = (

7
litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py
View file

@ -88,7 +88,10 @@ class AuthorizationCodeRefresher:
if token.refresh_token is None:
return None
server: Final = self._server_lookup(server_id)
if server is None or not server.token_url:
if server is None:
return None
token_url: Final = server.effective_token_url
if not token_url:
return None
try:
@ -106,7 +109,7 @@ class AuthorizationCodeRefresher:
"refresh_token": token.refresh_token,
**token_request.body,
}
body: Final = await self._token_endpoint(server.token_url, form, token_request.headers)
body: Final = await self._token_endpoint(token_url, form, token_request.headers)
if body is None:
return None
access_token: Final = body.get("access_token")

33
litellm/proxy/prompts/prompt_endpoints.py
View file

@ -1045,15 +1045,8 @@ async def delete_prompt(
raise HTTPException(status_code=500, detail=str(e))
def _reload_prompt_in_registry(
registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec
) -> PromptSpec:
"""Remove stale entry and re-initialize the prompt in the in-memory registry."""
if versioned_id in registry.IN_MEMORY_PROMPTS:
del registry.IN_MEMORY_PROMPTS[versioned_id]
if versioned_id in registry.prompt_id_to_custom_prompt:
del registry.prompt_id_to_custom_prompt[versioned_id]
initialized: Final = registry.initialize_prompt(prompt=updated_prompt_spec, config_file_path=None)
def _reload_prompt_in_registry(registry: "InMemoryPromptRegistry", updated_prompt_spec: PromptSpec) -> PromptSpec:
initialized: Final = registry.reload_prompt(prompt=updated_prompt_spec)
if initialized is None:
raise HTTPException(status_code=500, detail="Failed to patch prompt")
return initialized
@ -1146,25 +1139,15 @@ async def patch_prompt(
detail="Cannot update config prompts.",
)
# Use existing prompt from memory or build from DB row for field merging
if existing_prompt:
current_litellm_params = existing_prompt.litellm_params
current_prompt_info = existing_prompt.prompt_info
else:
current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row)
current_litellm_params = current_spec.litellm_params
current_prompt_info = current_spec.prompt_info
current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row)
# Update fields if provided
updated_litellm_params: Final = (
request.litellm_params if request.litellm_params is not None else current_litellm_params
request.litellm_params if request.litellm_params is not None else current_spec.litellm_params
)
updated_prompt_info: Final = request.prompt_info if request.prompt_info is not None else current_prompt_info
# Ensure we have valid litellm_params
if updated_litellm_params is None:
raise HTTPException(status_code=400, detail="litellm_params cannot be None")
updated_prompt_info: Final = (
request.prompt_info if request.prompt_info is not None else current_spec.prompt_info
)
# Build update data dict
update_data: Final[dict[str, str]] = {
@ -1188,7 +1171,7 @@ async def patch_prompt(
updated_prompt_spec: Final = create_versioned_prompt_spec(db_prompt=updated_prompt_db_entry)
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec)
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, updated_prompt_spec)
except HTTPException as e:
raise e

49
litellm/proxy/prompts/prompt_registry.py
View file

@ -118,7 +118,16 @@ class InMemoryPromptRegistry:
verbose_proxy_logger.debug("prompt_id already exists in IN_MEMORY_PROMPTS")
return self.IN_MEMORY_PROMPTS[prompt_id]
custom_prompt_callback: CustomPromptManagement | None = None
parsed_prompt, custom_prompt_callback = self._build_prompt_callback(prompt=prompt)
litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback)
# store references to the prompt in memory
self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt
self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback
return parsed_prompt
def _build_prompt_callback(self, prompt: PromptSpec) -> tuple[PromptSpec, CustomPromptManagement]:
litellm_params_data: Final = prompt.litellm_params
verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data)
@ -132,17 +141,17 @@ class InMemoryPromptRegistry:
raise ValueError("prompt_integration is required")
initializer: Final = prompt_initializer_registry.get(prompt_integration)
if initializer:
custom_prompt_callback = initializer(litellm_params, prompt)
if not isinstance(custom_prompt_callback, CustomPromptManagement):
raise ValueError(f"CustomPromptManagement is required, got {type(custom_prompt_callback)}")
litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback)
else:
if initializer is None:
raise ValueError(f"Unsupported prompt: {prompt_integration}")
custom_prompt_callback: Final = initializer(litellm_params, prompt)
if not isinstance(custom_prompt_callback, CustomPromptManagement):
raise ValueError( # noqa: TRY004 # prompt endpoints map ValueError to HTTP 400; keep the existing contract
f"CustomPromptManagement is required, got {type(custom_prompt_callback)}"
)
parsed_prompt: Final = PromptSpec(
prompt_id=prompt_id,
prompt_id=prompt.prompt_id,
litellm_params=litellm_params,
prompt_info=prompt.prompt_info or PromptInfo(prompt_type="config"),
created_at=prompt.created_at,
@ -151,13 +160,29 @@ class InMemoryPromptRegistry:
environment=prompt.environment,
created_by=prompt.created_by,
)
return parsed_prompt, custom_prompt_callback
# store references to the prompt in memory
self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt
self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback
def reload_prompt(self, prompt: PromptSpec) -> PromptSpec | None:
import litellm
parsed_prompt, new_callback = self._build_prompt_callback(prompt=prompt)
stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt.prompt_id, None)
self.IN_MEMORY_PROMPTS.pop(prompt.prompt_id, None)
if stale_callback is not None:
litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback)
litellm.logging_callback_manager.add_litellm_callback(new_callback)
self.IN_MEMORY_PROMPTS[prompt.prompt_id] = parsed_prompt
self.prompt_id_to_custom_prompt[prompt.prompt_id] = new_callback
return parsed_prompt
def sync_prompt_from_db(self, prompt: PromptSpec) -> PromptSpec | None:
existing: Final = self.IN_MEMORY_PROMPTS.get(prompt.prompt_id)
if existing is None:
return self.initialize_prompt(prompt=prompt)
if existing.litellm_params == prompt.litellm_params and existing.prompt_info == prompt.prompt_info:
return existing
return self.reload_prompt(prompt=prompt)
def get_prompt_by_id(self, prompt_id: str) -> PromptSpec | None:
"""
Get a prompt by its ID from memory

71
litellm/proxy/proxy_server.py
View file

@ -7257,12 +7257,40 @@ class ProxyConfig:
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.types.prompts.init_prompts import PromptSpec
def parse_row(db_prompt: object) -> PromptSpec | None:
try:
return self._get_prompt_spec_for_db_prompt(db_prompt=db_prompt)
except Exception as row_error: # noqa: BLE001 # a malformed row must not block syncing the remaining prompts
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to parse prompt row %s: %s",
getattr(db_prompt, "prompt_id", None),
row_error,
)
return None
try:
prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many()
for prompt in prompts_in_db:
# Convert DB object to dict and create versioned prompt_id
prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt)
IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec)
parsed_specs: Final[tuple[PromptSpec, ...]] = tuple(
spec for row in prompts_in_db if (spec := parse_row(row)) is not None
)
newest_spec_per_id: Final[Mapping[str, PromptSpec]] = MappingProxyType(
{
spec.prompt_id: spec
for spec in sorted(
parsed_specs,
key=lambda s: s.updated_at.timestamp() if s.updated_at else float("-inf"),
)
}
)
for prompt_spec in newest_spec_per_id.values():
try:
IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec)
except Exception as prompt_sync_error: # noqa: BLE001 # one poisoned row must not block syncing the remaining prompts
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to sync prompt %s: %s",
prompt_spec.prompt_id,
prompt_sync_error,
)
except Exception as e:
verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e)
@ -7497,11 +7525,9 @@ class ProxyConfig:
len(db_search_tools),
)
if llm_router is not None and search_tools:
if llm_router is not None:
await SearchAPIRouter.update_router_search_tools(router_instance=llm_router, search_tools=search_tools)
verbose_proxy_logger.info("Successfully loaded %s search tool(s) into router", len(search_tools))
elif llm_router is not None:
verbose_proxy_logger.debug("No search tools found in config or database, skipping router update")
else:
verbose_proxy_logger.debug(
"Router not initialized yet, search tools will be added when router is created"
@ -7512,6 +7538,26 @@ class ProxyConfig:
"litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - %s", e
)
async def reload_search_tools_from_db(self) -> None:
"""Refresh this worker's router from the search tools table.
Driven by the management endpoints so the worker that served the write is correct
immediately, and by the periodic job in store_model_in_db-off deployments. Gated the same
way as startup, so an admin who excluded search_tools from supported_db_objects opts out.
Serialized by MODEL_RECONCILE_LOCK for the reason add_deployment documents: the body is a
read-modify-write of the shared ``llm_router`` global, so two of them interleaving lets the
older snapshot's wholesale assignment land last and restore a tool the newer one deleted.
The lock belongs here rather than in _init_search_tools_in_db, which _init_non_llm_objects_in_db
already calls while holding it.
"""
if not self._should_load_db_object(object_type="search_tools"):
return
if prisma_client is None:
return
async with MODEL_RECONCILE_LOCK:
await self._init_search_tools_in_db(prisma_client=prisma_client)
@staticmethod
def _merge_config_and_db_search_tools(
config_search_tools: list[SearchToolTypedDict],
@ -9131,7 +9177,18 @@ class ProxyStartupEvent:
if store_model_in_db is not True:
await proxy_config.init_mcp_servers_from_db()
# Without this branch's own refresh, a UI-created search tool never reaches the router:
# the add_deployment job that carries it in store_model_in_db=True mode is not scheduled.
await proxy_config.reload_search_tools_from_db()
if prisma_client is not None:
scheduler.add_job(
proxy_config.reload_search_tools_from_db,
"interval",
seconds=config_reload_interval_seconds,
id="reload_search_tools_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
# DB-backed MCP servers are live objects in every mode, so the registry refresh that
# store_model_in_db=True deployments get via the add_deployment job must run here
# too; without it, a server whose OAuth discovery failed at startup is rebuilt only

28
litellm/proxy/search_endpoints/search_tool_management.py
View file

@ -51,6 +51,20 @@ def _convert_datetime_to_str(value: datetime | str | None) -> str | None:
TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]]
async def _refresh_router_search_tools() -> None:
"""Push the search tools table into this worker's router.
Best-effort: the row is already committed, so a refresh failure must not surface as a 500 and
push the caller into a retry that creates duplicates.
"""
from litellm.proxy.proxy_server import proxy_config
try:
await proxy_config.reload_search_tools_from_db()
except Exception as e: # noqa: BLE001 # the row is committed; no refresh failure may reach the caller
verbose_proxy_logger.exception("Search tool router refresh failed after a management write: %s", e)
async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable:
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import (
@ -305,8 +319,10 @@ async def create_search_tool(request: CreateSearchToolRequest):
search_tool=request.search_tool, prisma_client=prisma_client
)
await _refresh_router_search_tools()
verbose_proxy_logger.debug(
"Successfully added search tool '%s' to database. Router will be updated by the cron job.",
"Successfully added search tool '%s' to database.",
result.get("search_tool_name"),
)
@ -388,8 +404,10 @@ async def update_search_tool(search_tool_id: str, request: UpdateSearchToolReque
prisma_client=prisma_client,
)
await _refresh_router_search_tools()
verbose_proxy_logger.debug(
"Successfully updated search tool '%s' in database. Router will be updated by the cron job.",
"Successfully updated search tool '%s' in database.",
result.get("search_tool_name"),
)
@ -445,9 +463,9 @@ async def delete_search_tool(search_tool_id: str):
search_tool_id=search_tool_id, prisma_client=prisma_client
)
verbose_proxy_logger.debug(
"Successfully deleted search tool from database. Router will be updated by the cron job."
)
await _refresh_router_search_tools()
verbose_proxy_logger.debug("Successfully deleted search tool from database.")
return result
except HTTPException as e:

12
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -183,6 +183,18 @@ class MCPServer(BaseModel):
def __str__(self) -> str:
return self.__repr__()
@property
def effective_authorization_url(self) -> str | None:
return self.authorization_url or self.configured_authorization_url
@property
def effective_token_url(self) -> str | None:
return self.token_url or self.configured_token_url
@property
def effective_registration_url(self) -> str | None:
return self.registration_url or self.configured_registration_url
@property
def has_client_credentials(self) -> bool:
"""True if this server should use the OAuth2 client_credentials (M2M) flow.

147
model_prices_and_context_window.json
View file

@ -20150,7 +20150,7 @@
"gemini-3.5-flash-lite": {
"deprecation_date": "2027-07-21",
"cache_read_input_token_cost": 3e-08,
"cache_read_input_token_cost_flex": 2e-08,
"cache_read_input_token_cost_flex": 1.5e-08,
"cache_read_input_token_cost_priority": 5e-08,
"input_cost_per_token": 3e-07,
"input_cost_per_token_batches": 1.5e-07,
@ -20332,7 +20332,7 @@
"supports_image_size": false
},
"gemini-2.5-flash-preview-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost": 3e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-language-models",
@ -20377,10 +20377,53 @@
"google_maps_grounding_cost_per_query": 0.025,
"supports_image_size": false
},
"gemini-live-2.5-flash-native-audio": {
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "realtime",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supported_endpoints": [
"/vertex_ai/live"
],
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.035,
"search_context_size_medium": 0.035,
"search_context_size_high": 0.035
},
"gemini_native_audio": true
},
"gemini-live-2.5-flash-preview-native-audio-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
@ -20424,7 +20467,7 @@
"gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
@ -20469,7 +20512,7 @@
},
"gemini-2.5-flash-lite-preview-06-17": {
"deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost": 1e-08,
"input_cost_per_audio_token": 5e-07,
"input_cost_per_token": 1e-07,
"litellm_provider": "vertex_ai-language-models",
@ -21126,18 +21169,15 @@
},
"gemini-2.5-pro-preview-tts": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_above_200k_tokens": 2.5e-07,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"input_cost_per_token": 1e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_above_200k_tokens": 1.5e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview",
"output_cost_per_token": 2e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_modalities": [
"text"
],
@ -22062,7 +22102,7 @@
"supports_image_size": false
},
"gemini/gemini-2.5-flash-preview-09-2025": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost": 3e-08,
"deprecation_date": "2026-02-17",
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
@ -22111,7 +22151,7 @@
"supports_image_size": false
},
"gemini/gemini-flash-latest": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost": 3e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"litellm_provider": "gemini",
@ -22158,7 +22198,7 @@
"google_maps_grounding_cost_per_query": 0.025
},
"gemini/gemini-flash-lite-latest": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost": 1e-08,
"input_cost_per_audio_token": 3e-07,
"input_cost_per_token": 1e-07,
"litellm_provider": "gemini",
@ -22206,7 +22246,7 @@
},
"gemini/gemini-2.5-flash-lite-preview-06-17": {
"deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost": 1e-08,
"input_cost_per_audio_token": 5e-07,
"input_cost_per_token": 1e-07,
"litellm_provider": "gemini",
@ -22254,11 +22294,11 @@
"supports_image_size": false
},
"gemini/gemini-2.5-flash-preview-tts": {
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"mode": "audio_speech",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_token": 1e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/audio/speech"
],
@ -23212,19 +23252,16 @@
},
"gemini/gemini-2.5-pro-preview-tts": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_above_200k_tokens": 2.5e-07,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"input_cost_per_token": 1e-06,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_above_200k_tokens": 1.5e-05,
"output_cost_per_token": 2e-05,
"rpm": 10000,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview",
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_modalities": [
"text"
],
@ -42027,7 +42064,7 @@
"vertex_ai/gemini-3.5-flash-lite": {
"deprecation_date": "2027-07-21",
"cache_read_input_token_cost": 3e-08,
"cache_read_input_token_cost_flex": 2e-08,
"cache_read_input_token_cost_flex": 1.5e-08,
"cache_read_input_token_cost_priority": 5e-08,
"input_cost_per_token": 3e-07,
"input_cost_per_token_batches": 1.5e-07,
@ -48884,15 +48921,16 @@
}
},
"gemini-2.5-flash-native-audio-latest": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -48909,15 +48947,16 @@
"gemini_native_audio": true
},
"gemini-2.5-flash-native-audio-preview-09-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -48934,15 +48973,16 @@
"gemini_native_audio": true
},
"gemini-2.5-flash-native-audio-preview-12-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -48992,15 +49032,16 @@
"gemini_audio_only_live": true
},
"gemini/gemini-2.5-flash-native-audio-latest": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -49019,15 +49060,16 @@
"gemini_native_audio": true
},
"gemini/gemini-2.5-flash-native-audio-preview-09-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -49046,15 +49088,16 @@
"gemini_native_audio": true
},
"gemini/gemini-2.5-flash-native-audio-preview-12-2025": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
"input_cost_per_audio_token": 3e-06,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_audio_token": 1.2e-05,
"output_cost_per_token": 2e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/realtime"
],
@ -49123,11 +49166,11 @@
"rpm": 10
},
"gemini-2.5-flash-preview-tts": {
"input_cost_per_token": 3e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "gemini",
"mode": "audio_speech",
"output_cost_per_token": 2.5e-06,
"source": "https://ai.google.dev/pricing",
"output_cost_per_token": 1e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1/audio/speech"
]

47
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -3377,6 +3377,53 @@ def test_generic_cost_per_token_gemini_35_flash_lite(_local_model_cost_map):
assert completion_cost == pytest.approx(0.00125)
GEMINI_35_FLASH_LITE_TIER_RATES_BY_SURFACE = [
("gemini", None, 3e-07, 2.5e-06, 3e-08),
("gemini", "flex", 1.5e-07, 1.25e-06, 2e-08),
("gemini", "priority", 5.4e-07, 4.5e-06, 5e-08),
("vertex_ai", None, 3e-07, 2.5e-06, 3e-08),
("vertex_ai", "flex", 1.5e-07, 1.25e-06, 1.5e-08),
("vertex_ai", "priority", 5.4e-07, 4.5e-06, 5e-08),
]
@pytest.mark.parametrize(
"custom_llm_provider,service_tier,input_rate,output_rate,cache_read_rate",
GEMINI_35_FLASH_LITE_TIER_RATES_BY_SURFACE,
)
def test_gemini_35_flash_lite_service_tier_pricing(
custom_llm_provider, service_tier, input_rate, output_rate, cache_read_rate, _local_model_cost_map
):
"""Regression: Vertex publishes flash-lite flex context caching at $0.015/M while the
Gemini API publishes $0.02/M, so vertex_ai flex cache reads must bill 1.5e-08/token
instead of the 2e-08 the map used to carry, without disturbing the Gemini API rate."""
usage = Usage(
prompt_tokens=1_000,
completion_tokens=500,
total_tokens=1_500,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200, text_tokens=800),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="gemini-3.5-flash-lite",
usage=usage,
custom_llm_provider=custom_llm_provider,
service_tier=service_tier,
)
assert prompt_cost == pytest.approx(800 * input_rate + 200 * cache_read_rate, rel=1e-9)
assert completion_cost == pytest.approx(500 * output_rate, rel=1e-9)
def test_gemini_35_flash_lite_flex_cache_read_map_entries(_local_model_cost_map):
"""Each map entry carries its own surface's published flex cache-read rate: the bare
and vertex_ai keys are the Vertex surface at $0.015/M, the gemini key is the Gemini
API surface at $0.02/M."""
assert litellm.model_cost["gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 1.5e-08
assert litellm.model_cost["vertex_ai/gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 1.5e-08
assert litellm.model_cost["gemini/gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 2e-08
@pytest.mark.parametrize(
"service_tier,input_rate,cache_read_rate,cache_write_rate,output_rate",
[

49
tests/test_litellm/llms/gemini/test_cost_calculator.py
View file

@ -391,3 +391,52 @@ def test_map_traffic_type_to_service_tier(
assert (
_map_traffic_type_to_service_tier(traffic_type) == expected_service_tier
)
@pytest.mark.parametrize(
"model,custom_llm_provider,expected_cache_read_cost",
[
("gemini/gemini-flash-latest", "gemini", 3e-08),
("gemini/gemini-flash-lite-latest", "gemini", 1e-08),
("gemini/gemini-2.5-flash-preview-09-2025", "gemini", 3e-08),
("gemini/gemini-2.5-flash-lite-preview-06-17", "gemini", 1e-08),
("vertex_ai/gemini-2.5-flash-preview-09-2025", "vertex_ai", 3e-08),
("vertex_ai/gemini-2.5-flash-lite-preview-06-17", "vertex_ai", 1e-08),
],
)
def test_flash_alias_cache_read_is_ten_percent_of_input(
monkeypatch, model, custom_llm_provider, expected_cache_read_cost
):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
model_info = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
assert model_info["cache_read_input_token_cost"] == expected_cache_read_cost
assert model_info["cache_read_input_token_cost"] == pytest.approx(
0.10 * model_info["input_cost_per_token"]
)
@pytest.mark.parametrize(
"prefixed,bare",
[
("gemini/gemini-flash-latest", "gemini-flash-latest"),
("gemini/gemini-flash-lite-latest", "gemini-flash-lite-latest"),
],
)
def test_flash_latest_alias_spellings_price_identically(monkeypatch, prefixed, bare):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
prefixed_entry = litellm.model_cost[prefixed]
bare_entry = litellm.model_cost[bare]
for cost_key in (
"input_cost_per_token",
"output_cost_per_token",
"cache_read_input_token_cost",
):
assert prefixed_entry[cost_key] == bare_entry[cost_key]

19
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
View file

@ -579,3 +579,22 @@ def test_id_jag_honors_explicit_subject_token_type():
def test_id_jag_half_configured_defers_to_v1(server):
# A half-configured server must defer (None) rather than 500 at IdJagConfig construction.
assert to_server_spec(server) is None
def test_client_credentials_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
admin-entered value; the M2M spec must carry it so egress can mint."""
spec = to_server_spec(
_server(
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
url="https://up.example.com/mcp",
token_url=None,
configured_token_url="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
)
)
assert spec is not None
assert isinstance(spec.config, ClientCredentialsConfig)
assert spec.config.token_url == "https://idp.example.com/token"

25
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py
View file

@ -20,8 +20,10 @@ class _Server:
upstream_resource=None,
url=None,
server_id="srv",
configured_token_url=None,
):
self.token_url = token_url
self.configured_token_url = configured_token_url
self.client_id = client_id
self.client_secret = client_secret
self.token_endpoint_auth_method = token_endpoint_auth_method
@ -29,6 +31,10 @@ class _Server:
self.url = url
self.server_id = server_id
@property
def effective_token_url(self):
return self.token_url or self.configured_token_url
def _lookup(server):
return lambda server_id: server
@ -262,3 +268,22 @@ async def test_returned_scope_overrides_prior_when_present():
assert token is not None
assert token.scopes == ("read",) # a present scope replaces the prior grant
assert persisted[0][5] == ("read",)
@pytest.mark.asyncio
async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
admin-entered value; the refresh grant must POST there instead of silently failing."""
posted = []
refresher = _refresher(
server=_Server(token_url=None, configured_token_url="https://idp.example.com/token"),
body={"access_token": "new-at", "expires_in": 3600},
post_sink=posted,
)
token = await refresher.refresh(
"alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")
)
assert token is not None
assert token.access_token == "new-at"
assert posted[0][0] == "https://idp.example.com/token"

25
tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py
View file

@ -1337,3 +1337,28 @@ def test_mcp_oauth_token_identity_changes_when_only_upstream_resource_is_edited(
assert mcp_oauth_token_identity(set_to_explicit) == mcp_oauth_token_identity(
_identity_server(credentials={**creds, "upstream_resource": "api://audience-one"})
)
@pytest.mark.asyncio
async def test_refresh_user_oauth_token_uses_admin_entered_token_url_when_issuer_yield_empties_resolved(monkeypatch):
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
admin-entered value; the silent per-user refresh must POST there instead of bailing."""
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv-1",
name="test",
url="https://up.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
client_secret="csec",
token_url=None,
configured_token_url="https://idp.example.com/token",
)
result, captured = await _run_refresh(monkeypatch, server)
assert result is not None
assert captured["url"] == "https://idp.example.com/token"

446
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -79,24 +79,23 @@ def _resolved_oauth_metadata():
@pytest.mark.asyncio
async def test_authorize_resolves_cold_oauth_metadata():
async def test_authorize_resolves_cold_oauth_metadata(monkeypatch):
"""The route hands the registered server to the flow, whose deferred-discovery join resolves
the cold metadata; the redirect must land on the discovered authorization endpoint."""
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-lit-6255")
server = _unresolved_oauth_server()
global_mcp_server_manager.registry[server.server_id] = server
global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True)
request = _mock_callback_request("https://litellm.example.com/")
expected = MagicMock()
with (
patch.object(
global_mcp_server_manager,
"_discover_oauth_metadata_for_server",
new=AsyncMock(return_value=_resolved_oauth_metadata()),
) as discovery,
patch.object(discoverable_endpoints, "authorize_with_server", new=AsyncMock(return_value=expected)) as relay,
):
with patch.object(
global_mcp_server_manager,
"_discover_oauth_metadata_for_server",
new=AsyncMock(return_value=_resolved_oauth_metadata()),
) as discovery:
response = await discoverable_endpoints.authorize(
request=request,
client_id="client-id",
@ -105,12 +104,14 @@ async def test_authorize_resolves_cold_oauth_metadata():
)
discovery.assert_awaited_once_with(server)
assert relay.await_args.kwargs["mcp_server"].authorization_url == "https://idp.example.com/authorize"
assert response is expected
assert response.status_code == 307
assert response.headers["location"].startswith("https://idp.example.com/authorize")
@pytest.mark.asyncio
async def test_token_resolves_cold_oauth_metadata():
"""The route hands the registered server to the exchange, whose deferred-discovery join
resolves the cold metadata; the exchange must post to the discovered token endpoint."""
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
@ -118,7 +119,11 @@ async def test_token_resolves_cold_oauth_metadata():
global_mcp_server_manager.registry[server.server_id] = server
global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True)
request = _mock_callback_request("https://litellm.example.com/")
expected = MagicMock()
fake_http_response = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
with (
patch.object(
@ -127,8 +132,10 @@ async def test_token_resolves_cold_oauth_metadata():
new=AsyncMock(return_value=_resolved_oauth_metadata()),
) as discovery,
patch.object(
discoverable_endpoints, "exchange_token_with_server", new=AsyncMock(return_value=expected)
) as relay,
discoverable_endpoints,
"get_async_httpx_client",
new=lambda llm_provider: fake_http_client,
),
):
response = await discoverable_endpoints.token_endpoint(
request=request,
@ -139,20 +146,26 @@ async def test_token_resolves_cold_oauth_metadata():
)
discovery.assert_awaited_once_with(server)
assert relay.await_args.kwargs["mcp_server"].token_url == "https://idp.example.com/token"
assert response is expected
assert response.status_code == 200
assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/token"
@pytest.mark.asyncio
async def test_register_resolves_cold_oauth_metadata():
"""The route hands the registered server to the registration flow, whose deferred-discovery
join resolves the cold metadata; DCR must post to the discovered registration endpoint."""
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
server = _unresolved_oauth_server()
server = _unresolved_oauth_server().model_copy(update={"client_id": None})
global_mcp_server_manager.registry[server.server_id] = server
global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True)
request = _mock_callback_request("https://litellm.example.com/")
expected = MagicMock()
fake_http_response = MagicMock()
fake_http_response.json.return_value = {"client_id": "generated-client", "client_secret": "generated-secret"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
with (
patch.object(
@ -162,14 +175,135 @@ async def test_register_resolves_cold_oauth_metadata():
) as discovery,
patch.object(discoverable_endpoints, "_read_request_body", new=AsyncMock(return_value={})),
patch.object(
discoverable_endpoints, "register_client_with_server", new=AsyncMock(return_value=expected)
) as relay,
discoverable_endpoints,
"get_async_httpx_client",
new=lambda llm_provider: fake_http_client,
),
):
response = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_name)
discovery.assert_awaited_once_with(server)
assert relay.await_args.kwargs["mcp_server"].registration_url == "https://idp.example.com/register"
assert response is expected
assert response.status_code == 200
assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/register"
@pytest.mark.asyncio
async def test_register_route_bridge_missing_registration_url_joins_discovery():
"""A clientless DCR bridge whose authorize and token urls are admin-entered still relays
registration upstream: the flow must join deferred discovery for the missing registration
endpoint instead of short-circuiting to dummy credentials because authorization resolves."""
import json
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="bridge-partial-metadata",
name="bridge_partial_metadata",
server_name="bridge_partial_metadata",
alias="bridge_partial_metadata",
url="https://mcp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth_delegate,
dcr_bridge=True,
client_id=None,
authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/token",
)
global_mcp_server_manager.registry[server.server_id] = server
global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True)
request = _mock_callback_request("https://litellm.example.com/")
fake_http_response = MagicMock()
fake_http_response.json.return_value = {"client_id": "generated-client", "client_secret": "generated-secret"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
with (
patch.object( # test-quality-ok: innermost discovery seam on a module-global manager; the route-to-flow join under test stays real
global_mcp_server_manager,
"_discover_oauth_metadata_for_server",
new=AsyncMock(return_value=_resolved_oauth_metadata()),
) as discovery,
patch.object( # test-quality-ok: the MagicMock Request carries no body; this seam feeds the RFC 7591 redirect_uris
discoverable_endpoints,
"_read_request_body",
new=AsyncMock(return_value={"redirect_uris": ["https://client.example.com/cb"]}),
),
patch.object( # test-quality-ok: keeps the DCR POST off the network so its target URL can be asserted
discoverable_endpoints,
"get_async_httpx_client",
new=lambda llm_provider: fake_http_client,
),
):
response = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_name)
discovery.assert_awaited_once_with(server)
assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/register"
assert fake_http_client.post.await_args.kwargs["json"]["redirect_uris"] == ["https://client.example.com/cb"]
assert response.status_code == 200
assert json.loads(response.body.decode("utf-8"))["client_id"] == "generated-client"
@pytest.mark.asyncio
async def test_token_route_bridge_missing_registration_url_joins_discovery():
"""A clientless DCR bridge rebuilt with an admin-entered token url but without its discovered
registration endpoint must rejoin discovery at the exchange: the relay-vs-callback arm hinges
on the registration url, so skipping discovery would swap the client's own redirect_uri for
the gateway callback and the upstream would reject the code."""
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="bridge-partial-token-metadata",
name="bridge_partial_token_metadata",
server_name="bridge_partial_token_metadata",
alias="bridge_partial_token_metadata",
url="https://mcp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
dcr_bridge=True,
client_id=None,
authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/token",
)
global_mcp_server_manager.registry[server.server_id] = server
global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True)
request = _mock_callback_request("https://litellm.example.com/")
fake_http_response = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
with (
patch.object( # test-quality-ok: innermost discovery seam on a module-global manager; the route-to-exchange join under test stays real
global_mcp_server_manager,
"_discover_oauth_metadata_for_server",
new=AsyncMock(return_value=_resolved_oauth_metadata()),
) as discovery,
patch.object( # test-quality-ok: keeps the upstream token POST off the network so its redirect_uri arm can be asserted
discoverable_endpoints,
"get_async_httpx_client",
new=lambda llm_provider: fake_http_client,
),
):
response = await discoverable_endpoints.token_endpoint(
request=request,
grant_type="authorization_code",
code="upstream-code",
redirect_uri="https://client.example.com/cb",
client_id="dcr-client-id",
mcp_server_name=server.server_name,
)
discovery.assert_awaited_once_with(server)
assert response.status_code == 200
assert fake_http_client.post.await_args.kwargs["data"]["redirect_uri"] == "https://client.example.com/cb"
@pytest.fixture
@ -8878,6 +9012,272 @@ async def test_authorize_wall_names_the_issuer_for_anchored_servers():
assert "idp.example.com" not in detail_text
@pytest.mark.asyncio
async def test_authorize_uses_admin_entered_github_oauth_urls_after_issuer_yield(monkeypatch):
"""GitHub MCP servers store Authorization URL and Token URL on the row. 1.99 can empty
the resolved authorization_url when a leftover issuer is treated as a pin (RFC 8414
yield). The UI authorize must still redirect to the admin-entered GitHub authorize URL
instead of 400ing that discovery against api.githubcopilot.com failed."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
authorize_with_server,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="ecac50c4-8eca-438a-af80-9bdebadafc69",
name="github_mcp",
alias="github_mcp",
server_name="github_mcp",
url="https://api.githubcopilot.com/mcp/",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
client_id="github-app-client",
authorization_url=None,
token_url=None,
issuer="https://github.com",
issuer_is_anchored=True,
configured_authorization_url="https://github.com/login/oauth/authorize",
configured_token_url="https://github.com/login/oauth/access_token",
)
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-lit-6255")
response = await authorize_with_server(
request=mock_request,
mcp_server=server,
client_id="github-app-client",
redirect_uri="http://127.0.0.1:60108/callback",
state="state123",
)
assert response.status_code == 307
assert "https://github.com/login/oauth/authorize" in response.headers["location"]
assert "client_id=github-app-client" in response.headers["location"]
def test_oauth_endpoints_count_admin_entered_urls_as_resolved():
"""A leftover issuer empties the resolved authorize/token fields but must not keep the
server on the deferred-discovery retry path when the admin already stored those URLs."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_oauth_endpoints_unresolved,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="github-configured",
name="github_mcp",
url="https://api.githubcopilot.com/mcp/",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
authorization_url=None,
token_url=None,
configured_authorization_url="https://github.com/login/oauth/authorize",
configured_token_url="https://github.com/login/oauth/access_token",
)
assert _oauth_endpoints_unresolved(server) is False
@pytest.mark.asyncio
async def test_token_exchange_with_configured_token_url_never_joins_discovery(monkeypatch):
"""A server can hold an admin-entered Token URL while its Authorization URL is absent. The
token exchange must post to that stored endpoint without awaiting deferred discovery, which
can 503 against an unreachable issuer even though nothing it resolves is needed here."""
from litellm.proxy._experimental.mcp_server import (
discoverable_endpoints,
mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
exchange_token_with_server,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="token-url-only",
name="token_url_only",
server_name="token_url_only",
alias="token_url_only",
url="https://mcp.example.com/mcp/",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
client_secret="cs",
authorization_url=None,
token_url=None,
issuer="https://idp.example.com",
issuer_is_anchored=True,
configured_token_url="https://idp.example.com/oauth/token",
)
async def fail_discovery(_srv):
raise AssertionError("the exchange joined deferred discovery despite a stored token url")
monkeypatch.setattr(
mcp_server_manager.global_mcp_server_manager,
"ensure_oauth_metadata_discovered",
fail_discovery,
)
fake_http_response = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
monkeypatch.setattr(
discoverable_endpoints,
"get_async_httpx_client",
lambda llm_provider: fake_http_client,
)
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
response = await exchange_token_with_server(
request=mock_request,
mcp_server=server,
grant_type="authorization_code",
code="upstream-code",
redirect_uri="http://127.0.0.1:3000/cb",
client_id="cid",
client_secret=None,
code_verifier=None,
)
assert response.status_code == 200
assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/oauth/token"
@pytest.mark.asyncio
async def test_root_token_route_with_configured_token_url_never_joins_discovery(monkeypatch):
"""A root POST /token that falls back to the sole OAuth2 server must reach the exchange's
endpoint-gated discovery join instead of awaiting full discovery at the route: with the
token url admin-entered, a failing or slow discovery must not turn the exchange into a 503."""
from litellm.proxy._experimental.mcp_server import (
discoverable_endpoints,
mcp_server_manager,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
manager = mcp_server_manager.global_mcp_server_manager
server = MCPServer(
server_id="sole-token-url-only",
name="sole_token_url_only",
server_name="sole_token_url_only",
alias="sole_token_url_only",
url="https://mcp.example.com/mcp/",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
client_secret="cs",
issuer="https://idp.example.com",
issuer_is_anchored=True,
configured_token_url="https://idp.example.com/oauth/token",
)
saved_registry = dict(manager.registry)
manager.registry.clear()
manager.registry[server.server_id] = server
manager._set_oauth_discovery_deferred(server.server_id, True)
async def fail_discovery(_srv):
raise AssertionError("the root token route joined deferred discovery despite a stored token url")
monkeypatch.setattr(manager, "ensure_oauth_metadata_discovered", fail_discovery)
fake_http_response = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
monkeypatch.setattr(
discoverable_endpoints,
"get_async_httpx_client",
lambda llm_provider: fake_http_client,
)
request = _mock_callback_request("https://litellm.example.com/")
try:
response = await discoverable_endpoints.token_endpoint(
request=request,
grant_type="authorization_code",
code="upstream-code",
redirect_uri="http://127.0.0.1:3000/cb",
client_id="unregistered-dcr-client",
)
finally:
manager.registry.clear()
manager.registry.update(saved_registry)
assert response.status_code == 200
assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/oauth/token"
@pytest.mark.asyncio
async def test_bridge_authorize_relays_with_registration_url_resolved_by_deferred_discovery(monkeypatch):
"""When deferred discovery resolves a DCR-bridge server during the authorize request, the
relay-vs-short-circuit call must read the resolved server: a client that registered itself
through the front door keeps its own redirect binding instead of being routed through the
gateway callback the upstream never granted it."""
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
authorize_with_server,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="bridge-deferred",
name="bridge_deferred",
server_name="bridge_deferred",
alias="bridge_deferred",
url="https://mcp.example.com/mcp/",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
dcr_bridge=True,
authorization_url=None,
token_url=None,
registration_url=None,
)
resolved = server.model_copy(
update={
"authorization_url": "https://idp.example.com/oauth/authorize",
"token_url": "https://idp.example.com/oauth/token",
"registration_url": "https://idp.example.com/oauth/register",
}
)
async def resolve_discovery(_srv):
return resolved
monkeypatch.setattr(
mcp_server_manager.global_mcp_server_manager,
"ensure_oauth_metadata_discovered",
resolve_discovery,
)
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
response = await authorize_with_server(
request=mock_request,
mcp_server=server,
client_id="front-door-client",
redirect_uri="http://127.0.0.1:60110/client-callback",
state="state456",
code_challenge="E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM",
code_challenge_method="S256",
)
assert response.status_code == 307
location = response.headers["location"]
assert location.startswith("https://idp.example.com/oauth/authorize")
assert "redirect_uri=http%3A%2F%2F127.0.0.1%3A60110%2Fclient-callback" in location
def test_passthrough_authorization_code_round_trips_and_rejects_hostile_input():
"""The passthrough gateway code seals and recovers the ephemeral DCR client and upstream code,
and is total over hostile input: a raw upstream code opens to None, and a tampered or

19
tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py
View file

@ -392,3 +392,22 @@ async def test_invalidate_clears_every_identity_for_a_server():
assert refetched == "tok-after-invalidate"
assert mock_client.post.call_count == 3
@pytest.mark.asyncio
async def test_m2m_mint_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
admin-entered value; the client_credentials mint must POST there instead of raising."""
server = _server(token_url=None, configured_token_url="https://auth.example.com/token")
cache = MCPOAuth2TokenCache()
mock_client = AsyncMock()
mock_client.post.return_value = _token_response("m2m-token-configured")
with patch(
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
return_value=mock_client,
):
result = await cache.async_get_token(server)
assert result == "m2m-token-configured"
assert mock_client.post.call_args[0][0] == "https://auth.example.com/token"

156
tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py
View file

@ -992,3 +992,159 @@ async def test_list_search_tools_reports_a_missing_real_team_as_404():
assert response.status_code == 404
assert "search_tools" not in response.json()
# ---------------------------------------------------------------------------
# Router sync on management writes (LIT-3379)
#
# The proxy resolves prisma_client / proxy_config / llm_router from
# litellm.proxy.proxy_server module globals at call time and reaches its DB layer through a
# module-level registry singleton, so there is no constructor or parameter to inject through.
# Patching those globals is the only seam that exercises the endpoint end to end.
# ---------------------------------------------------------------------------
def _search_tool_row(name: str, provider: str = "tavily") -> dict:
return {
"search_tool_id": f"{name}-id",
"search_tool_name": name,
"litellm_params": {"search_provider": provider, "api_key": "sk-test"},
"search_tool_info": {"description": name},
}
def _fake_registry(db_rows: list) -> MagicMock:
"""A registry singleton whose writes land in db_rows, so the refresh reads back real state."""
async def _add(search_tool, **_):
row = _search_tool_row(
search_tool["search_tool_name"],
provider=search_tool.get("litellm_params", {}).get("search_provider", "tavily"),
)
db_rows.append(row)
return row
async def _update(search_tool_id, search_tool, **_):
row = _search_tool_row(
search_tool["search_tool_name"],
provider=search_tool.get("litellm_params", {}).get("search_provider", "tavily"),
)
db_rows[:] = [row if existing["search_tool_id"] == search_tool_id else existing for existing in db_rows]
return row
async def _delete(search_tool_id, **_):
db_rows[:] = [existing for existing in db_rows if existing["search_tool_id"] != search_tool_id]
return {"message": "deleted", "search_tool_name": search_tool_id}
async def _get_by_id(search_tool_id, **_):
return next((row for row in db_rows if row["search_tool_id"] == search_tool_id), None)
registry = MagicMock()
registry.add_search_tool_to_db = AsyncMock(side_effect=_add)
registry.update_search_tool_in_db = AsyncMock(side_effect=_update)
registry.delete_search_tool_from_db = AsyncMock(side_effect=_delete)
registry.get_search_tool_by_id_from_db = AsyncMock(side_effect=_get_by_id)
return registry
@contextlib.contextmanager
def _live_router_and_db(db_rows: list):
"""Drive the endpoints against a real ProxyConfig so the router refresh actually runs."""
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
proxy_config.update_config_state({})
fake_router = MagicMock()
fake_router.search_tools = list(db_rows)
with contextlib.ExitStack() as stack:
stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", MagicMock())) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", proxy_config)) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", fake_router)) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
_fake_registry(db_rows),
)
)
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
AsyncMock(side_effect=lambda **_: list(db_rows)),
)
)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
)
try:
yield fake_router
finally:
app.dependency_overrides.pop(user_api_key_auth, None)
@pytest.mark.asyncio
async def test_create_search_tool_reaches_the_router_before_the_response():
"""A UI-created tool must be usable immediately, not only after the next config reload tick."""
with _live_router_and_db([]) as fake_router:
response = TestClient(app).post(
"/search_tools",
json={
"search_tool": {
"search_tool_name": "tavily-search",
"litellm_params": {"search_provider": "tavily"},
}
},
)
assert response.status_code == 200
assert [tool["search_tool_name"] for tool in fake_router.search_tools] == ["tavily-search"]
@pytest.mark.asyncio
async def test_update_search_tool_reaches_the_router_before_the_response():
with _live_router_and_db([_search_tool_row("tavily-search", provider="tavily")]) as fake_router:
response = TestClient(app).put(
"/search_tools/tavily-search-id",
json={
"search_tool": {
"search_tool_name": "tavily-search",
"litellm_params": {"search_provider": "exa_ai"},
}
},
)
assert response.status_code == 200
assert fake_router.search_tools[0]["litellm_params"]["search_provider"] == "exa_ai"
@pytest.mark.asyncio
async def test_delete_search_tool_removes_it_from_the_router():
"""Deleting the last tool must clear the router; the old empty-list guard left it live."""
with _live_router_and_db([_search_tool_row("tavily-search")]) as fake_router:
response = TestClient(app).delete("/search_tools/tavily-search-id")
assert response.status_code == 200
assert fake_router.search_tools == []
@pytest.mark.asyncio
async def test_create_search_tool_survives_a_failing_router_refresh():
"""The row is already committed, so a refresh failure must not turn into a 500."""
with _live_router_and_db([]):
with patch( # test-quality-ok: forcing the refresh to fail needs the refresh itself replaced
"litellm.proxy.proxy_server.ProxyConfig.reload_search_tools_from_db",
AsyncMock(side_effect=RuntimeError("registry boom")),
):
response = TestClient(app).post(
"/search_tools",
json={
"search_tool": {
"search_tool_name": "tavily-search",
"litellm_params": {"search_provider": "tavily"},
}
},
)
assert response.status_code == 200
assert response.json()["search_tool_name"] == "tavily-search"

76
tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py
View file

@ -1,4 +1,5 @@
import json
import pytest
from unittest.mock import MagicMock, AsyncMock, patch
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
@ -9,6 +10,27 @@ from litellm.types.prompts.init_prompts import (
)
def _db_row(content: str) -> MagicMock:
row = MagicMock()
row.id = "row-1"
row.version = 1
row.model_dump.return_value = {
"prompt_id": "test_prompt",
"version": 1,
"environment": "development",
"created_by": None,
"litellm_params": {
"prompt_id": "test_prompt",
"prompt_integration": "dotprompt",
"prompt_data": {"content": content, "metadata": {}},
},
"prompt_info": {"prompt_type": "db"},
"created_at": None,
"updated_at": None,
}
return row
@pytest.mark.asyncio
async def test_delete_prompt_success():
"""
@ -209,9 +231,7 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404():
api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN
)
target_row = MagicMock()
target_row.id = "row-1"
target_row.version = 1
target_row = _db_row("Begin every reply with AHOY")
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(
@ -249,6 +269,48 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404():
)
@pytest.mark.asyncio
async def test_patch_prompt_merges_unsent_fields_from_db_row_not_stale_memory():
from litellm.proxy.prompts.prompt_endpoints import PatchPromptRequest, patch_prompt
mock_user_auth = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN)
db_row = _db_row("Begin every reply with HOWDY")
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row])
mock_prisma_client.db.litellm_prompttable.update = AsyncMock(return_value=db_row)
stale_in_memory = PromptSpec(
prompt_id="test_prompt.v1",
litellm_params=PromptLiteLLMParams(
prompt_id="test_prompt",
prompt_integration="dotprompt",
prompt_data={"content": "Begin every reply with AHOY", "metadata": {}},
),
prompt_info=PromptInfo(prompt_type="db"),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch( # test-quality-ok: stubs the collaborator so the test pins what the endpoint writes and reloads
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
) as mock_registry,
):
mock_registry.get_prompt_by_id.return_value = stale_in_memory
mock_registry.reload_prompt.side_effect = lambda prompt: prompt
response = await patch_prompt(
prompt_id="test_prompt",
request=PatchPromptRequest(prompt_info=PromptInfo(prompt_type="db")),
user_api_key_dict=mock_user_auth,
)
written_params = json.loads(mock_prisma_client.db.litellm_prompttable.update.call_args.kwargs["data"]["litellm_params"])
assert written_params["prompt_data"]["content"] == "Begin every reply with HOWDY"
reloaded_spec = mock_registry.reload_prompt.call_args.kwargs["prompt"]
assert reloaded_spec.prompt_id == "test_prompt.v1"
assert reloaded_spec.litellm_params.prompt_data["content"] == "Begin every reply with HOWDY"
assert response.litellm_params.prompt_data["content"] == "Begin every reply with HOWDY"
def test_is_ambiguous_keyed_prompt_data_shapes():
from litellm.proxy.prompts.prompt_endpoints import is_ambiguous_keyed_prompt_data
@ -358,6 +420,14 @@ async def test_patch_prompt_info_only_keeps_legacy_keyed_row_patchable():
target_row = MagicMock()
target_row.id = "row-1"
target_row.version = 1
target_row.model_dump.return_value = {
"prompt_id": "agent-prompt",
"version": 1,
"environment": "production",
"created_by": None,
"litellm_params": legacy_params.model_dump_json(),
"prompt_info": PromptInfo(prompt_type="db", environment="production").model_dump_json(),
}
updated_row = MagicMock()
updated_row.model_dump.return_value = {
"prompt_id": "agent-prompt",

90
tests/test_litellm/proxy/prompts/test_prompt_registry.py Normal file
View file

@ -0,0 +1,90 @@
import pytest
import litellm
from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry
from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec
def _db_prompt_spec(content: str) -> PromptSpec:
return PromptSpec(
prompt_id="greeting.v1",
litellm_params=PromptLiteLLMParams(
prompt_id="greeting",
prompt_integration="dotprompt",
prompt_data={"content": content, "metadata": {}},
),
prompt_info=PromptInfo(prompt_type="db"),
)
def _served_content(registry: InMemoryPromptRegistry) -> str:
callback = registry.get_prompt_callback_by_id("greeting.v1")
assert callback is not None
return callback.prompt_manager.get_prompt("greeting").content
@pytest.fixture
def isolated_callbacks(monkeypatch: pytest.MonkeyPatch) -> list:
monkeypatch.setattr(litellm, "callbacks", [])
return litellm.callbacks
def test_sync_prompt_from_db_reloads_row_edited_elsewhere(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY"))
stale_callback = registry.get_prompt_callback_by_id("greeting.v1")
assert _served_content(registry) == "begin every reply with AHOY"
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY"))
assert _served_content(registry) == "begin every reply with HOWDY"
assert registry.get_prompt_by_id("greeting.v1").litellm_params.prompt_data["content"] == "begin every reply with HOWDY"
assert stale_callback not in isolated_callbacks
assert isolated_callbacks == [registry.get_prompt_callback_by_id("greeting.v1")]
def test_sync_prompt_from_db_keeps_unchanged_row_in_place(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY"))
first_callback = registry.get_prompt_callback_by_id("greeting.v1")
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY"))
assert registry.get_prompt_callback_by_id("greeting.v1") is first_callback
assert isolated_callbacks == [first_callback]
def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY"))
stale_callback = registry.get_prompt_callback_by_id("greeting.v1")
reloaded = registry.reload_prompt(prompt=_db_prompt_spec("begin every reply with HOWDY"))
assert reloaded is not None
assert _served_content(registry) == "begin every reply with HOWDY"
assert stale_callback not in isolated_callbacks
assert len(isolated_callbacks) == 1
def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY"))
old_callback = registry.get_prompt_callback_by_id("greeting.v1")
broken = PromptSpec(
prompt_id="greeting.v1",
litellm_params=PromptLiteLLMParams(
prompt_id="greeting",
prompt_integration="does_not_exist",
prompt_data={"content": "begin every reply with HOWDY", "metadata": {}},
),
prompt_info=PromptInfo(prompt_type="db"),
)
with pytest.raises(ValueError, match="Unsupported prompt"):
registry.reload_prompt(prompt=broken)
assert registry.get_prompt_callback_by_id("greeting.v1") is old_callback
assert _served_content(registry) == "begin every reply with AHOY"
assert isolated_callbacks == [old_callback]

100
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -1253,26 +1253,114 @@ async def test_ProxyConfig__init_search_tools_in_db_loads_merged_tools(monkeypat
@pytest.mark.asyncio
async def test_ProxyConfig__init_search_tools_in_db_skips_empty_router_update(monkeypatch):
async def test_ProxyConfig__init_search_tools_in_db_clears_router_when_last_tool_is_deleted(monkeypatch):
"""Deleting the last search tool must clear the router, not leave the tool live in memory."""
from litellm.proxy import proxy_server
from litellm.router_utils.search_api_router import SearchAPIRouter
pc = ProxyConfig()
pc.update_config_state({})
fake_router = MagicMock()
fake_router.search_tools = [{"search_tool_name": "deleted-search", "litellm_params": {}}]
mock_get_db_tools = AsyncMock(return_value=[])
mock_update_router = AsyncMock()
monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
mock_get_db_tools,
)
monkeypatch.setattr(SearchAPIRouter, "update_router_search_tools", mock_update_router)
await pc._init_search_tools_in_db(prisma_client=MagicMock())
mock_get_db_tools.assert_awaited_once()
mock_update_router.assert_not_awaited()
assert fake_router.search_tools == []
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_refreshes_router(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
await pc.reload_search_tools_from_db()
mock_init.assert_awaited_once()
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_honors_supported_db_objects(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
await pc.reload_search_tools_from_db()
mock_init.assert_not_awaited()
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_serializes_overlapping_refreshes(monkeypatch):
"""An older snapshot must not land last and restore a tool a newer refresh deleted."""
import asyncio
from litellm.proxy import proxy_server
pc = ProxyConfig()
pc.update_config_state({})
fake_router = MagicMock()
fake_router.search_tools = []
stale_read_started = asyncio.Event()
fresh_write_committed = asyncio.Event()
snapshots = iter(
(
[{"search_tool_name": "doomed-search", "litellm_params": {}}],
[],
)
)
async def _read_db(**_):
snapshot = next(snapshots)
if not stale_read_started.is_set():
stale_read_started.set()
await fresh_write_committed.wait()
return snapshot
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
_read_db,
)
stale = asyncio.create_task(pc.reload_search_tools_from_db())
await stale_read_started.wait()
deleter = asyncio.create_task(pc.reload_search_tools_from_db())
await asyncio.sleep(0)
fresh_write_committed.set()
await asyncio.gather(stale, deleter)
assert fake_router.search_tools == []
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_noops_without_prisma(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", None)
await pc.reload_search_tools_from_db()
mock_init.assert_not_awaited()
# ---------------------------------------------------------------------------

141
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -11351,6 +11351,147 @@ async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_re
assert not GUARDRAIL_RECONCILE_LOCK.locked()
@pytest.mark.asyncio
async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeypatch):
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.proxy.proxy_server import ProxyConfig
monkeypatch.setattr(litellm, "callbacks", [])
def db_row(content: str) -> MagicMock:
row = MagicMock()
row.model_dump.return_value = {
"prompt_id": "greeting_sync",
"version": 1,
"environment": "development",
"created_by": None,
"litellm_params": json.dumps(
{
"prompt_id": "greeting_sync",
"prompt_integration": "dotprompt",
"prompt_data": {"content": content, "metadata": {}},
}
),
"prompt_info": json.dumps({"prompt_type": "db"}),
"created_at": None,
"updated_at": None,
}
return row
def served_content() -> str:
callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1")
assert callback is not None
return callback.prompt_manager.get_prompt("greeting_sync").content
prisma_client = MagicMock()
try:
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with AHOY")])
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert served_content() == "Begin every reply with AHOY"
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with HOWDY")])
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert served_content() == "Begin every reply with HOWDY"
assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1")]
finally:
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_sync")
@pytest.mark.asyncio
async def test_init_prompts_in_db_syncs_remaining_rows_when_one_row_fails(monkeypatch):
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.proxy.proxy_server import ProxyConfig
monkeypatch.setattr(litellm, "callbacks", [])
def db_row(prompt_id: str, integration: str) -> MagicMock:
row = MagicMock()
row.model_dump.return_value = {
"prompt_id": prompt_id,
"version": 1,
"environment": "development",
"created_by": None,
"litellm_params": json.dumps(
{
"prompt_id": prompt_id,
"prompt_integration": integration,
"prompt_data": {"content": "Begin every reply with AHOY", "metadata": {}},
}
),
"prompt_info": json.dumps({"prompt_type": "db"}),
"created_at": None,
"updated_at": None,
}
return row
prisma_client = MagicMock()
try:
prisma_client.db.litellm_prompttable.find_many = AsyncMock(
return_value=[db_row("broken_sync", "does_not_exist"), db_row("healthy_sync", "dotprompt")]
)
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("broken_sync.v1") is None
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1") is not None
assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1")]
finally:
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("healthy_sync")
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("broken_sync")
@pytest.mark.asyncio
async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collide_on_a_versioned_id(monkeypatch):
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.proxy.proxy_server import ProxyConfig
monkeypatch.setattr(litellm, "callbacks", [])
def db_row(environment: str, content: str, updated_at: datetime) -> MagicMock:
row = MagicMock()
row.model_dump.return_value = {
"prompt_id": "greeting_env",
"version": 1,
"environment": environment,
"created_by": None,
"litellm_params": json.dumps(
{
"prompt_id": "greeting_env",
"prompt_integration": "dotprompt",
"prompt_data": {"content": content, "metadata": {}},
}
),
"prompt_info": json.dumps({"prompt_type": "db"}),
"created_at": None,
"updated_at": updated_at,
}
return row
freshly_patched = db_row(
"production", "Begin every reply with HOWDY", datetime(2026, 8, 26, 12, 0, tzinfo=timezone.utc)
)
stale_sibling = db_row(
"development", "Begin every reply with AHOY", datetime(2026, 8, 26, 11, 0, tzinfo=timezone.utc)
)
prisma_client = MagicMock()
try:
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[freshly_patched, stale_sibling])
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
first_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1")
assert first_callback is not None
assert first_callback.prompt_manager.get_prompt("greeting_env").content == "Begin every reply with HOWDY"
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") is first_callback
assert litellm.callbacks == [first_callback]
finally:
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_env")
class TestEmbeddingsFailureHookRequestData:
@pytest.mark.asyncio
async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):

98
tests/test_litellm/test_cost_calculator.py
View file

@ -4159,3 +4159,101 @@ def test_every_one_hour_cache_write_rate_is_double_its_input_rate():
}
assert deviations == {}
def test_gemini_live_native_audio_ga_realtime_cost(_local_model_cost_map: None) -> None:
"""Regression for https://github.com/BerriAI/litellm/issues/31087."""
from litellm.types.utils import CompletionTokensDetailsWrapper
results: OpenAIRealtimeStreamList = [
{"type": "session.created", "session": {"model": "gemini-live-2.5-flash-native-audio"}},
]
combined_usage_object = Usage(
prompt_tokens=8,
completion_tokens=25,
total_tokens=33,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=8, audio_tokens=0),
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=2, audio_tokens=23),
)
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage_object,
custom_llm_provider="vertex_ai",
litellm_model_name="vertex_ai/gemini-live-2.5-flash-native-audio",
)
expected_cost = 8 * 5e-07 + 2 * 2e-06 + 23 * 1.2e-05
assert cost == pytest.approx(expected_cost, rel=1e-9)
@pytest.mark.parametrize(
"priceless_entry",
[
{"litellm_provider": "vertex_ai", "mode": "realtime"},
{
"litellm_provider": "vertex_ai",
"mode": "realtime",
"input_cost_per_token": None,
"output_cost_per_token": None,
"input_cost_per_audio_token": None,
},
],
ids=["registered_without_price_fields", "registered_with_none_valued_price_fields"],
)
def test_realtime_priceless_deployment_entry_falls_through_to_priced_model(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, priceless_entry: dict
) -> None:
"""Regression for https://github.com/BerriAI/litellm/issues/31087 (router-registered priceless entries)."""
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/some-unmapped-live-model",
priceless_entry,
)
priced_model = "vertex_ai/gemini-live-2.5-flash-preview-native-audio-09-2025"
priced_entry = litellm.model_cost["gemini-live-2.5-flash-preview-native-audio-09-2025"]
results: OpenAIRealtimeStreamList = [
{"type": "session.created", "session": {"model": "some-unmapped-live-model"}},
]
combined_usage_object = Usage(prompt_tokens=8, completion_tokens=25, total_tokens=33)
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage_object,
custom_llm_provider="vertex_ai",
litellm_model_name=priced_model,
)
expected_cost = 8 * priced_entry["input_cost_per_token"] + 25 * priced_entry["output_cost_per_token"]
assert cost == pytest.approx(expected_cost, rel=1e-9)
assert cost > 0
def test_realtime_explicitly_free_session_model_still_bills_zero(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/free-live-model",
{
"litellm_provider": "vertex_ai",
"mode": "realtime",
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
},
)
results: OpenAIRealtimeStreamList = [
{"type": "session.created", "session": {"model": "free-live-model"}},
]
combined_usage_object = Usage(prompt_tokens=8, completion_tokens=25, total_tokens=33)
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage_object,
custom_llm_provider="vertex_ai",
litellm_model_name="vertex_ai/gemini-live-2.5-flash-preview-native-audio-09-2025",
)
assert cost == 0.0

151
tests/test_litellm/test_gemini_tts_native_audio_pricing.py Normal file
View file

@ -0,0 +1,151 @@
import json
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import pytest
import litellm
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage
REPO_ROOT: Final = Path(__file__).parents[2]
MAIN_PATH: Final = REPO_ROOT / "model_prices_and_context_window.json"
BACKUP_PATH: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json"
FLASH_TTS_KEYS: Final = ("gemini-2.5-flash-preview-tts", "gemini/gemini-2.5-flash-preview-tts")
PRO_TTS_KEYS: Final = ("gemini-2.5-pro-preview-tts", "gemini/gemini-2.5-pro-preview-tts")
NATIVE_AUDIO_KEYS: Final = tuple(
f"{prefix}gemini-2.5-flash-native-audio-{suffix}"
for prefix in ("", "gemini/")
for suffix in ("latest", "preview-09-2025", "preview-12-2025")
)
LIVE_NATIVE_AUDIO_KEYS: Final = (
"gemini-live-2.5-flash-preview-native-audio-09-2025",
"gemini/gemini-live-2.5-flash-preview-native-audio-09-2025",
)
FLASH_TTS_INPUT: Final = 5e-07
FLASH_TTS_AUDIO_OUTPUT: Final = 1e-05
PRO_TTS_INPUT: Final = 1e-06
PRO_TTS_AUDIO_OUTPUT: Final = 2e-05
NATIVE_AUDIO_TEXT_INPUT: Final = 5e-07
NATIVE_AUDIO_AUDIO_INPUT: Final = 3e-06
NATIVE_AUDIO_TEXT_OUTPUT: Final = 2e-06
NATIVE_AUDIO_AUDIO_OUTPUT: Final = 1.2e-05
PUBLISHED_RATES: Final = {
**{
key: {"input_cost_per_token": FLASH_TTS_INPUT, "output_cost_per_token": FLASH_TTS_AUDIO_OUTPUT}
for key in FLASH_TTS_KEYS
},
**{
key: {"input_cost_per_token": PRO_TTS_INPUT, "output_cost_per_token": PRO_TTS_AUDIO_OUTPUT}
for key in PRO_TTS_KEYS
},
**{
key: {
"input_cost_per_token": NATIVE_AUDIO_TEXT_INPUT,
"input_cost_per_audio_token": NATIVE_AUDIO_AUDIO_INPUT,
"output_cost_per_token": NATIVE_AUDIO_TEXT_OUTPUT,
"output_cost_per_audio_token": NATIVE_AUDIO_AUDIO_OUTPUT,
}
for key in (*NATIVE_AUDIO_KEYS, *LIVE_NATIVE_AUDIO_KEYS)
},
}
ALL_KEYS: Final = tuple(PUBLISHED_RATES)
NATIVE_AUDIO_BILLING_CASES: Final = (
*((key, "gemini") for key in NATIVE_AUDIO_KEYS),
("gemini-live-2.5-flash-preview-native-audio-09-2025", "vertex_ai"),
("gemini/gemini-live-2.5-flash-preview-native-audio-09-2025", "gemini"),
)
LONG_CONTEXT_TIER_FIELDS: Final = (
"input_cost_per_token_above_200k_tokens",
"output_cost_per_token_above_200k_tokens",
"cache_read_input_token_cost_above_200k_tokens",
)
def _load(path: Path) -> dict[str, dict[str, object]]:
with open(path, encoding="utf-8") as f:
return json.load(f)
@pytest.fixture
def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
litellm.get_model_info.cache_clear()
yield
litellm.get_model_info.cache_clear()
@pytest.mark.parametrize("model", ALL_KEYS)
@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup"))
def test_published_rates_are_registered(model: str, path: Path):
info = _load(path)[model]
for field, value in PUBLISHED_RATES[model].items():
assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}"
@pytest.mark.parametrize("model", PRO_TTS_KEYS)
@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup"))
def test_pro_tts_has_no_long_context_tier(model: str, path: Path):
info = _load(path)[model]
for field in LONG_CONTEXT_TIER_FIELDS:
assert field not in info, f"{model} has {field} but Google publishes one flat TTS rate"
@pytest.mark.parametrize("model", ALL_KEYS)
def test_backup_matches_main(model: str):
assert _load(BACKUP_PATH)[model] == _load(MAIN_PATH)[model]
@pytest.mark.parametrize(
("model", "provider", "input_rate", "audio_output_rate"),
(
("gemini-2.5-flash-preview-tts", "gemini", FLASH_TTS_INPUT, FLASH_TTS_AUDIO_OUTPUT),
("gemini-2.5-pro-preview-tts", "gemini", PRO_TTS_INPUT, PRO_TTS_AUDIO_OUTPUT),
("gemini-2.5-pro-preview-tts", "vertex_ai", PRO_TTS_INPUT, PRO_TTS_AUDIO_OUTPUT),
),
)
def test_tts_audio_output_is_billed_at_the_audio_rate(
model: str, provider: str, input_rate: float, audio_output_rate: float, local_model_cost_map
):
usage: Final = Usage(
prompt_tokens=9,
completion_tokens=49,
total_tokens=58,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=9),
completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=49, text_tokens=0),
)
prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider)
assert prompt_cost == pytest.approx(9 * input_rate)
assert completion_cost == pytest.approx(49 * audio_output_rate)
@pytest.mark.parametrize("model, provider", NATIVE_AUDIO_BILLING_CASES)
def test_native_audio_output_is_billed_at_the_audio_rate(model: str, provider: str, local_model_cost_map):
usage: Final = Usage(
prompt_tokens=377,
completion_tokens=84,
total_tokens=461,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=377),
completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=48, reasoning_tokens=36, text_tokens=0),
)
prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider)
assert prompt_cost == pytest.approx(377 * NATIVE_AUDIO_TEXT_INPUT)
assert completion_cost == pytest.approx(48 * NATIVE_AUDIO_AUDIO_OUTPUT + 36 * NATIVE_AUDIO_TEXT_OUTPUT)
@pytest.mark.parametrize("model, provider", NATIVE_AUDIO_BILLING_CASES)
def test_native_audio_input_is_billed_at_the_audio_rate(model: str, provider: str, local_model_cost_map):
usage: Final = Usage(
prompt_tokens=1000,
completion_tokens=0,
total_tokens=1000,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100, audio_tokens=900),
)
prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider)
assert prompt_cost == pytest.approx(100 * NATIVE_AUDIO_TEXT_INPUT + 900 * NATIVE_AUDIO_AUDIO_INPUT)