ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #24149 from michelligabriele/fix/kill-orphaned-prisma-engine-on-reconnect

Browse Source 
fix(proxy): kill orphaned prisma engine subprocess on failed disconnect
This commit is contained in:
Krish Dholakia 2026-03-19 18:12:58 -07:00 committed by GitHub
commit 75386d4cb9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 1051 additions and 245 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
.circleci/config.yml
View File

@ -42,7 +42,7 @@ commands:
"pydantic==2.11.0" "mcp==1.25.0" "requests-mock>=1.12.1" \
"responses==0.25.7" "pytest-xdist==3.6.1" "pytest-timeout==2.2.0" \
"pytest-cov==5.0.0" "semantic_router==0.1.10" "fastapi-offline==1.7.3" \
"a2a"
"a2a" "parameterized>=0.9.0"
- setup_litellm_enterprise_pip
- save_cache:
paths:
@ -1115,7 +1115,7 @@ jobs:
for dir in "${IGNORE_DIRS[@]}"; do
IGNORE_ARGS="$IGNORE_ARGS --ignore=$dir"
done
python -m pytest -v tests/llm_translation $IGNORE_ARGS --junitxml=test-results/junit.xml --durations=20 -n 8 --timeout=120 --timeout_method=thread
python -m pytest -v tests/llm_translation $IGNORE_ARGS --junitxml=test-results/junit.xml --durations=20 -n 8 --timeout=120 --timeout_method=thread --retries 2 --retry-delay 5
no_output_timeout: 15m
# Store test results
@ -1331,7 +1331,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/unified_google_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
python -m pytest -vv tests/unified_google_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 --retries 3 --retry-delay 5
no_output_timeout: 15m
- run:
name: Rename the coverage files

3
ci_cd/security_scans.sh
View File

@ -163,6 +163,9 @@ run_grype_scans() {
"CVE-2026-25639" # axios - full fix requires 1.x major version bump; pinned to >=0.30.2 to clear other axios CVEs, upgrade to 1.x in follow-up
"CVE-2026-2297" # Python 3.13 SourcelessFileLoader audit hook bypass - no fix available in base image
"GHSA-qffp-2rhf-9h96" # tar hardlink path traversal - from nodejs_wheel bundled npm, not used in application runtime code
"CVE-2026-2673" # OpenSSL 3.6.1 TLS 1.3 key exchange group negotiation issue - no fix available yet
"CVE-2026-3644" # Python 3.13 vulnerability - no fix available in base image
"CVE-2026-4224" # Python 3.13 Expat parser stack overflow in ElementDeclHandler - no fix available in base image
)
# Build JSON array of allowlisted CVE IDs for jq

78
docs/my-website/blog/guardrail_logging_secret_exposure_incident/index.md Normal file
View File

@ -0,0 +1,78 @@
---
slug: guardrail-logging-secret-exposure-incident
title: "Incident Report: Guardrail logging exposed secret headers in spend logs and traces"
date: 2026-03-18T10:00:00
authors:
- litellm
tags: [incident-report, security, guardrails]
hide_table_of_contents: false
---
**Date:** March 18, 2026
**Duration:** Unknown
**Severity:** High
**Status:** Resolved
## Summary
When a custom guardrail returned the full LiteLLM request/data dictionary, the guardrail response logged by LiteLLM could include `secret_fields.raw_headers`, including plaintext `Authorization` headers containing API keys or other credentials.
This information could then propagate to logging and observability surfaces that consume guardrail metadata, including:
- **Spend logs in the LiteLLM UI:** visible to admins with access to spend-log data
- **OpenTelemetry traces:** visible to anyone with access to the relevant telemetry backend
LLM calls, proxy routing, and provider execution were not blocked by this bug. The impact was exposure of sensitive request headers in observability and logging paths.
{/* truncate */}
---
## Background
LiteLLM keeps internal request data (including request headers) for use during the call. That data is not meant to be written to logs or telemetry.
When custom guardrails run, their outcomes are logged so they can appear in spend logs, OpenTelemetry traces, and other observability backends. If a guardrail returned the full request payload instead of a minimal result, that internal request data could be included in what was logged. Before the fix, the guardrail logging path did not strip that data before sending it to those systems.
```mermaid
flowchart TD
inboundRequest["1. Incoming proxy request"] --> storeSecrets["2. Store internal request data"]
storeSecrets --> guardrailRuns["3. Custom guardrail runs"]
guardrailRuns --> fullDataReturn["4. Guardrail returns full request payload"]
fullDataReturn --> loggingBuild["5. Build guardrail log payload"]
loggingBuild --> spendLogs["6a. Persist to spend logs / UI"]
loggingBuild --> otelTraces["6b. Attach to OTEL guardrail spans"]
```
---
## Root Cause
The root cause was incomplete sanitization in the guardrail logging path. When building the payload that gets sent to spend logs and traces, LiteLLM prepared guardrail responses for logging but did not strip internal request data (such as headers) from them. If a guardrail returned a response that included that data, it was passed through to the logging and observability systems unchanged.
---
## Impact
This issue required all of the following:
1. A custom guardrail returned the full LiteLLM request/data dictionary, or another response object containing `secret_fields`.
2. LiteLLM logged that guardrail response through the standard guardrail logging path.
3. An operator, admin, or telemetry consumer had access to the resulting logs or traces.
When those conditions were met, sensitive values could become visible through:
- **Spend logs / UI responses:** guardrail metadata could be included in spend-log payloads rendered in the admin UI.
- **OpenTelemetry traces:** `guardrail_response` could be written as a span attribute on guardrail spans.
- **Other downstream observability backends:** any integration consuming the same guardrail metadata could receive the leaked values.
This was a logging and telemetry exposure bug. It did not let callers bypass auth, access other tenants directly, or change model behavior, but it could expose plaintext credentials to people with access to those observability systems.
---
## Guidance For Users
- Upgrade to LiteLLM 1.82.3+.
- If you operated custom guardrails that return the full request/data dict, review whether spend logs or telemetry traces were retained during the affected period.
- Rotate any credentials that may have appeared in `Authorization` or other forwarded request headers in those systems.
- Apply least-privilege access controls to spend-log views and telemetry backends that may contain request-derived metadata.

1
docs/my-website/docs/proxy/config_settings.md
View File

@ -902,6 +902,7 @@ router_settings:
| OTEL_SERVICE_NAME | Service name identifier for OpenTelemetry
| OTEL_TRACER_NAME | Tracer name for OpenTelemetry tracing
| OTEL_LOGS_EXPORTER | Exporter type for OpenTelemetry logs (e.g., console)
| OTEL_IGNORE_CONTEXT_PROPAGATION | When true, ignore parent span context propagation in OpenTelemetry callbacks
| PAGERDUTY_API_KEY | API key for PagerDuty Alerting
| PANW_PRISMA_AIRS_API_KEY | API key for PANW Prisma AIRS service
| PANW_PRISMA_AIRS_API_BASE | Base URL for PANW Prisma AIRS service

2
litellm/cost_calculator.py
View File

@ -757,7 +757,7 @@ def _map_traffic_type_to_service_tier(traffic_type: Optional[str]) -> Optional[s
"""
if traffic_type is None:
return None
service_tier = _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER.get(traffic_type.upper())
service_tier = _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER.get(str(traffic_type).upper())
return service_tier

10
litellm/integrations/datadog/datadog.py
View File

@ -291,7 +291,7 @@ class DataDogLogger(
dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=safe_dumps(message_payload),
service=get_datadog_service(),
@ -442,7 +442,7 @@ class DataDogLogger(
verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload)
dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(standard_logging_object=standard_logging_object),
ddtags=",".join(get_datadog_tags(standard_logging_object=standard_logging_object)),
hostname=get_datadog_hostname(),
message=json_payload,
service=get_datadog_service(),
@ -545,7 +545,7 @@ class DataDogLogger(
_dd_message_str = safe_dumps(_payload_dict)
_dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=_dd_message_str,
service=get_datadog_service(),
@ -587,7 +587,7 @@ class DataDogLogger(
_dd_message_str = safe_dumps(_payload_dict)
_dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=_dd_message_str,
service=get_datadog_service(),
@ -678,7 +678,7 @@ class DataDogLogger(
dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=json_payload,
service=get_datadog_service(),

11
litellm/integrations/datadog/datadog_handler.py
View File

@ -38,8 +38,13 @@ def get_datadog_pod_name() -> str:
def get_datadog_tags(
standard_logging_object: Optional[StandardLoggingPayload] = None,
) -> str:
"""Build Datadog tags string used by multiple integrations."""
) -> List[str]:
"""Build Datadog tags as a list of individual tag strings.
Returns a list of "key:value" strings suitable for Datadog LLM Observability
(which expects tags as an array). For Datadog Logs API (ddtags), join with
comma: ",".join(get_datadog_tags(...)).
"""
base_tags = {
"env": get_datadog_env(),
@ -66,4 +71,4 @@ def get_datadog_tags(
if team_tag:
tags.append(f"team:{team_tag}")
return ",".join(tags)
return tags

4
litellm/integrations/datadog/datadog_llm_obs.py
View File

@ -203,7 +203,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
type="span",
attributes=DDSpanAttributes(
ml_app=get_datadog_service(),
tags=[get_datadog_tags()],
tags=get_datadog_tags(),
spans=self.log_queue,
),
),
@ -315,7 +315,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
duration=int((end_time - start_time).total_seconds() * 1e9),
metrics=metrics,
status="error" if error_info else "ok",
tags=[get_datadog_tags(standard_logging_object=standard_logging_payload)],
tags=get_datadog_tags(standard_logging_object=standard_logging_payload),
)
apm_trace_id = self._get_apm_trace_id()

26
litellm/litellm_core_utils/streaming_handler.py
View File

@ -1893,15 +1893,23 @@ class CustomStreamWrapper:
"usage",
getattr(complete_streaming_response, "usage"),
)
self.cache_streaming_response(
processed_chunk=complete_streaming_response.model_copy(
try:
_cache_copy = complete_streaming_response.model_copy(
deep=True
),
)
_log_copy = complete_streaming_response.model_copy(
deep=True
)
except RuntimeError:
_cache_copy = complete_streaming_response.model_copy()
_log_copy = complete_streaming_response.model_copy()
self.cache_streaming_response(
processed_chunk=_cache_copy,
cache_hit=cache_hit,
)
executor.submit(
self.logging_obj.success_handler,
complete_streaming_response.model_copy(deep=True),
_log_copy,
None,
None,
cache_hit,
@ -2113,11 +2121,15 @@ class CustomStreamWrapper:
"usage",
getattr(complete_streaming_response, "usage"),
)
try:
_copy = complete_streaming_response.model_copy(
deep=True
)
except RuntimeError:
_copy = complete_streaming_response.model_copy()
asyncio.create_task(
self.async_cache_streaming_response(
processed_chunk=complete_streaming_response.model_copy(
deep=True
),
processed_chunk=_copy,
cache_hit=cache_hit,
)
)

72
litellm/model_prices_and_context_window_backup.json
View File

@ -4462,6 +4462,78 @@
"supports_vision": true,
"supports_web_search": true
},
"azure/gpt-5.4-mini": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-5.4-nano": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-image-1": {
"cache_read_input_image_token_cost": 2.5e-06,
"cache_read_input_token_cost": 1.25e-06,

56
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View File

@ -208,26 +208,52 @@ async def exchange_token_with_server(
client_id: str,
client_secret: Optional[str],
code_verifier: Optional[str],
refresh_token: Optional[str] = None,
scope: Optional[str] = None,
):
if grant_type != "authorization_code":
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:
raise HTTPException(status_code=400, detail="MCP server token url is not set")
proxy_base_url = get_request_base_url(request)
token_data = {
"grant_type": "authorization_code",
"client_id": mcp_server.client_id if mcp_server.client_id else client_id,
"client_secret": mcp_server.client_secret
if mcp_server.client_secret
else client_secret,
"code": code,
"redirect_uri": f"{proxy_base_url}/callback",
}
resolved_client_id = mcp_server.client_id if mcp_server.client_id else client_id
resolved_client_secret = (
mcp_server.client_secret if mcp_server.client_secret else client_secret
)
if code_verifier:
token_data["code_verifier"] = code_verifier
if grant_type == "refresh_token":
if not refresh_token:
raise HTTPException(
status_code=400,
detail="refresh_token is required for refresh_token grant",
)
token_data: dict = {
"grant_type": "refresh_token",
"refresh_token": refresh_token,
"client_id": resolved_client_id,
}
if resolved_client_secret is not None:
token_data["client_secret"] = resolved_client_secret
if scope:
token_data["scope"] = scope
else:
if not code:
raise HTTPException(
status_code=400,
detail="code is required for authorization_code grant",
)
proxy_base_url = get_request_base_url(request)
token_data = {
"grant_type": "authorization_code",
"client_id": resolved_client_id,
"code": code,
"redirect_uri": f"{proxy_base_url}/callback",
}
if resolved_client_secret is not None:
token_data["client_secret"] = resolved_client_secret
if code_verifier:
token_data["code_verifier"] = code_verifier
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
response = await async_client.post(
@ -375,6 +401,8 @@ async def token_endpoint(
client_id: str = Form(...),
client_secret: Optional[str] = Form(None),
code_verifier: str = Form(None),
refresh_token: Optional[str] = Form(None),
scope: Optional[str] = Form(None),
mcp_server_name: Optional[str] = None,
):
"""
@ -408,6 +436,8 @@ async def token_endpoint(
client_id=client_id,
client_secret=client_secret,
code_verifier=code_verifier,
refresh_token=refresh_token,
scope=scope,
)

147
litellm/proxy/auth/auth_checks.py
View File

@ -29,6 +29,7 @@ from litellm.constants import (
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.proxy._types import (
RBAC_ROLES,
@ -407,18 +408,19 @@ async def common_checks( # noqa: PLR0915
# 2. If team can call model
if _model and team_object:
if not await can_team_access_model(
model=_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=valid_token.team_model_aliases if valid_token else None,
):
raise ProxyException(
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
)
with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"):
if not await can_team_access_model(
model=_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=valid_token.team_model_aliases if valid_token else None,
):
raise ProxyException(
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
)
# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
if valid_token is not None and valid_token.agent_id:
@ -443,54 +445,60 @@ async def common_checks( # noqa: PLR0915
## 2.1 If user can call model (if personal key)
if _model and team_object is None and user_object is not None:
await can_user_call_model(
model=_model,
llm_router=llm_router,
user_object=user_object,
)
with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"):
await can_user_call_model(
model=_model,
llm_router=llm_router,
user_object=user_object,
)
# 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget)
await _run_project_checks(
project_object=project_object,
_model=_model,
llm_router=llm_router,
skip_budget_checks=skip_budget_checks,
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.common_checks.run_project_checks"):
await _run_project_checks(
project_object=project_object,
_model=_model,
llm_router=llm_router,
skip_budget_checks=skip_budget_checks,
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
# If this is a free model, skip all budget checks
if not skip_budget_checks:
# 3. If team is in budget
await _team_max_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.team_max_budget_check"):
await _team_max_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 3.0.5. If team is over soft budget (alert only, doesn't block)
await _team_soft_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"):
await _team_soft_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 3.1. If organization is in budget
await _organization_max_budget_check(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.common_checks.organization_max_budget_check"):
await _organization_max_budget_check(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await _tag_max_budget_check(
request_body=request_body,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.tag_max_budget_check"):
await _tag_max_budget_check(
request_body=request_body,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 4. If user is in budget
## 4.1 check personal budget, if personal key
@ -508,14 +516,15 @@ async def common_checks( # noqa: PLR0915
)
## 4.2 check team member budget, if team key
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_budget"):
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
if (
@ -554,19 +563,21 @@ async def common_checks( # noqa: PLR0915
)
# 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
await vector_store_access_check(
request_body=request_body,
team_object=team_object,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.vector_store_access_check"):
await vector_store_access_check(
request_body=request_body,
team_object=team_object,
valid_token=valid_token,
)
# 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path)
await check_tools_allowlist(
request_body=request_body,
valid_token=valid_token,
team_object=team_object,
route=route,
)
with tracer.trace("litellm.proxy.auth.common_checks.check_tools_allowlist"):
await check_tools_allowlist(
request_body=request_body,
valid_token=valid_token,
team_object=team_object,
route=route,
)
return True

264
litellm/proxy/auth/user_api_key_auth.py
View File

@ -548,13 +548,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
custom_auth_api_key: bool = False
try:
# get the request body
await pre_db_read_auth_checks(
request_data=request_data,
request=request,
route=route,
)
with tracer.trace("litellm.proxy.auth.pre_db_read_auth_checks"):
await pre_db_read_auth_checks(
request_data=request_data,
request=request,
route=route,
)
pass_through_endpoints: Optional[List[dict]] = general_settings.get(
"pass_through_endpoints", None
)
@ -588,9 +587,10 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
### USER-DEFINED AUTH FUNCTION ###
if enterprise_custom_auth is not None:
response = await enterprise_custom_auth(
request=request, api_key=api_key, user_custom_auth=user_custom_auth
)
with tracer.trace("litellm.proxy.auth.enterprise_custom_auth"):
response = await enterprise_custom_auth(
request=request, api_key=api_key, user_custom_auth=user_custom_auth
)
if response is not None and isinstance(response, UserAPIKeyAuth):
validated = UserAPIKeyAuth.model_validate(response)
validated = await _run_post_custom_auth_checks(
@ -706,18 +706,19 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
# Fall through to virtual key checks
if do_standard_jwt_auth:
result = await JWTAuthManager.auth_builder(
request_data=request_data,
general_settings=general_settings,
api_key=api_key,
jwt_handler=jwt_handler,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=_safe_get_request_headers(request),
)
with tracer.trace("litellm.proxy.auth.jwt_auth_builder"):
result = await JWTAuthManager.auth_builder(
request_data=request_data,
general_settings=general_settings,
api_key=api_key,
jwt_handler=jwt_handler,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=_safe_get_request_headers(request),
)
is_proxy_admin = result["is_proxy_admin"]
team_id = result["team_id"]
@ -909,15 +910,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
try:
end_user_params["end_user_id"] = end_user_id
# get end-user object
_end_user_object = await get_end_user_object(
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
with tracer.trace("litellm.proxy.auth.get_end_user_object"):
_end_user_object = await get_end_user_object(
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
if _end_user_object is not None:
end_user_params[
"allowed_model_region"
@ -960,14 +961,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if valid_token is None:
## Check CACHE
try:
valid_token = await get_key_object(
hashed_token=hash_token(api_key),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
check_cache_only=True,
)
with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"):
valid_token = await get_key_object(
hashed_token=hash_token(api_key),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
check_cache_only=True,
)
except Exception:
verbose_logger.debug("api key not found in cache.")
valid_token = None
@ -1139,13 +1141,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
api_key = hash_token(token=api_key)
try:
valid_token = await get_key_object(
hashed_token=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.get_key_object_from_db"):
valid_token = await get_key_object(
hashed_token=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except ProxyException as e:
if e.code == 401 or e.code == "401":
e.message = "Authentication Error, Invalid proxy server token passed. Received API Key = {}, Key Hash (Token) ={}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format(
@ -1233,14 +1236,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
# Check 2. If user_id for this token is in budget - done in common_checks()
if valid_token.user_id is not None:
try:
user_obj = await get_user_object(
user_id=valid_token.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.get_user_object"):
user_obj = await get_user_object(
user_id=valid_token.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except Exception as e:
verbose_logger.debug(
"litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - {}".format(
@ -1329,71 +1333,73 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
if not skip_budget_checks:
# Check 4. Token Spend is under budget
if RouteChecks.is_llm_api_route(route=route):
await _virtual_key_max_budget_check(
with tracer.trace("litellm.proxy.auth.budget_checks"):
# Check 4. Token Spend is under budget
if RouteChecks.is_llm_api_route(route=route):
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Max Budget Alert Check
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Max Budget Alert Check
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 6. Soft Budget Check
await _virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Token Model Spend is under Model budget
max_budget_per_model = valid_token.model_max_budget
current_model = request_data.get("model", None)
if (
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and prisma_client is not None
and current_model is not None
and valid_token.token is not None
):
## GET THE SPEND FOR THIS MODEL
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
# Check 6. Soft Budget Check
await _virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5b. End-user model max budget
end_user_mmb = valid_token.end_user_model_max_budget
if (
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
# Check 5. Token Model Spend is under Model budget
max_budget_per_model = valid_token.model_max_budget
current_model = request_data.get("model", None)
if (
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and prisma_client is not None
and current_model is not None
and valid_token.token is not None
):
## GET THE SPEND FOR THIS MODEL
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
)
# Check 5b. End-user model max budget
end_user_mmb = valid_token.end_user_model_max_budget
if (
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
# Check 6: Additional Common Checks across jwt + key auth
if valid_token.team_id is not None:
try:
_team_obj = await get_team_object(
team_id=valid_token.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.get_team_object"):
_team_obj = await get_team_object(
team_id=valid_token.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except HTTPException:
_team_obj = LiteLLM_TeamTableCachedObj(
team_id=valid_token.team_id,
@ -1431,11 +1437,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
litellm.max_budget > 0 and prisma_client is not None
): # user set proxy max budget
cache_key = "{}:spend".format(litellm_proxy_admin_name)
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
cache_key=cache_key,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"):
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
cache_key=cache_key,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
if global_proxy_spend is not None:
call_info = CallInfo(
@ -1452,21 +1459,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_info=call_info,
)
)
_ = await common_checks(
request=request,
request_body=request_data,
team_object=_team_obj,
user_object=user_obj,
end_user_object=_end_user_object,
general_settings=general_settings,
global_proxy_spend=global_proxy_spend,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=skip_budget_checks,
project_object=_project_obj,
)
with tracer.trace("litellm.proxy.auth.common_checks"):
_ = await common_checks(
request=request,
request_body=request_data,
team_object=_team_obj,
user_object=user_obj,
end_user_object=_end_user_object,
general_settings=general_settings,
global_proxy_spend=global_proxy_spend,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=skip_budget_checks,
project_object=_project_obj,
)
# Token passed all checks
if valid_token is None:
raise HTTPException(401, detail="Invalid API key")

4
litellm/proxy/common_request_processing.py
View File

@ -1512,7 +1512,9 @@ class ProxyBaseLLMRequestProcessing:
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=(
_litellm_logging_obj.litellm_call_id if _litellm_logging_obj else None
_litellm_logging_obj.litellm_call_id
if _litellm_logging_obj
else self.data.get("litellm_call_id")
),
model_id=model_id,
version=version,

44
litellm/proxy/db/prisma_client.py
View File

@ -5,6 +5,7 @@ This file contains the PrismaWrapper class, which is used to wrap the Prisma cli
import asyncio
import os
import random
import signal
import subprocess
import time
import urllib
@ -45,6 +46,46 @@ class PrismaWrapper:
self._reconnection_lock = asyncio.Lock()
self._last_refresh_time: Optional[datetime] = None
def _get_engine_pid(self) -> int:
"""Get the PID of the current Prisma engine subprocess, or 0 if unavailable."""
try:
engine = self._original_prisma._engine
process = getattr(engine, "process", None) if engine is not None else None
if process is not None:
return process.pid
except (AttributeError, TypeError):
pass
return 0
@staticmethod
async def _kill_engine_process(pid: int) -> None:
"""Force-kill an orphaned engine subprocess to prevent DB connection pool leaks.
Called when disconnect() fails and the old engine process may still be
holding open connections. Sends SIGTERM for graceful shutdown, waits
briefly, then SIGKILL as a backstop.
"""
if pid <= 0:
return
try:
os.kill(pid, signal.SIGTERM)
except (ProcessLookupError, PermissionError, OSError):
return # Already dead or inaccessible
verbose_proxy_logger.warning(
"Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.",
pid,
)
# Brief wait for graceful shutdown, then force-kill
await asyncio.sleep(0.5)
try:
os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM))
verbose_proxy_logger.warning(
"Sent SIGKILL to prisma-query-engine PID %s (did not exit after SIGTERM).",
pid,
)
except (ProcessLookupError, PermissionError, OSError):
pass # Exited after SIGTERM — expected
def _extract_token_from_db_url(self, db_url: Optional[str]) -> Optional[str]:
"""
Extract the token (password) from the DATABASE_URL.
@ -179,10 +220,13 @@ class PrismaWrapper:
"""Disconnect and reconnect the Prisma client with a new database URL."""
from prisma import Prisma # type: ignore
old_engine_pid = self._get_engine_pid()
try:
await self._original_prisma.disconnect()
except Exception as e:
verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}")
await self._kill_engine_process(old_engine_pid)
if http_client is not None:
self._original_prisma = Prisma(http=http_client)

4
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View File

@ -1399,6 +1399,8 @@ if MCP_AVAILABLE:
client_id: Optional[str] = Form(None),
client_secret: Optional[str] = Form(None),
code_verifier: Optional[str] = Form(None),
refresh_token: Optional[str] = Form(None),
scope: Optional[str] = Form(None),
):
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
resolved_client_id = mcp_server.client_id or client_id or ""
@ -1422,6 +1424,8 @@ if MCP_AVAILABLE:
client_id=resolved_client_id,
client_secret=client_secret,
code_verifier=code_verifier,
refresh_token=refresh_token,
scope=scope,
)
@router.post(

7
litellm/proxy/utils.py
View File

@ -1917,6 +1917,7 @@ class ProxyLogging:
original_exception,
traceback.format_exc(),
),
daemon=True,
).start()
async def post_call_success_hook(
@ -4005,13 +4006,15 @@ class PrismaClient:
)
async def _do_direct_reconnect() -> None:
old_pid = self._get_engine_pid()
try:
await self.db.disconnect()
except Exception as disconnect_err:
verbose_proxy_logger.debug(
"Prisma DB disconnect before reconnect failed (ignored): %s",
verbose_proxy_logger.warning(
"Prisma DB disconnect before reconnect failed: %s",
disconnect_err,
)
await PrismaWrapper._kill_engine_process(old_pid)
await self.db.connect()
await self.db.query_raw("SELECT 1")

224
model_prices_and_context_window.json
View File

@ -4462,6 +4462,78 @@
"supports_vision": true,
"supports_web_search": true
},
"azure/gpt-5.4-mini": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-5.4-nano": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-image-1": {
"cache_read_input_image_token_cost": 2.5e-06,
"cache_read_input_token_cost": 1.25e-06,
@ -37032,5 +37104,157 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"volcengine/doubao-seed-2-0-pro-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 4.6e-07,
"output_cost_per_token": 2.3e-06,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 7e-07,
"output_cost_per_token": 3.5e-06,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 7e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-lite-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 8.7e-08,
"output_cost_per_token": 5.2e-07,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 1.3e-07,
"output_cost_per_token": 7.8e-07,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 2.6e-07,
"output_cost_per_token": 1.6e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-mini-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 2.9e-08,
"output_cost_per_token": 2.9e-07,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 5.8e-08,
"output_cost_per_token": 5.8e-07,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.2e-07,
"output_cost_per_token": 1.2e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-code-preview-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 4.6e-07,
"output_cost_per_token": 2.3e-06,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 7e-07,
"output_cost_per_token": 3.5e-06,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 7e-06,
"range": [
128000.0,
256000.0
]
}
]
}
}

18
tests/logging_callback_tests/test_datadog.py
View File

@ -593,7 +593,7 @@ def test_datadog_static_methods():
# Test tags format with default values
assert (
"env:unknown,service:litellm-server,version:unknown,HOSTNAME:"
in get_datadog_tags()
in ",".join(get_datadog_tags())
)
# Test with custom environment variables
@ -631,7 +631,7 @@ def test_datadog_static_methods():
# Test tags format with custom values
expected_custom_tags = "env:production,service:custom-service,version:1.0.0,HOSTNAME:test-host,POD_NAME:pod-123"
print("DataDogLogger._get_datadog_tags()", get_datadog_tags())
assert get_datadog_tags() == expected_custom_tags
assert ",".join(get_datadog_tags()) == expected_custom_tags
@pytest.mark.asyncio
@ -672,11 +672,11 @@ def test_get_datadog_tags():
"""Test the _get_datadog_tags static method with various inputs"""
# Test with no standard_logging_object and default env vars
base_tags = get_datadog_tags()
assert "env:" in base_tags
assert "service:" in base_tags
assert "version:" in base_tags
assert "POD_NAME:" in base_tags
assert "HOSTNAME:" in base_tags
assert any("env:" in t for t in base_tags)
assert any("service:" in t for t in base_tags)
assert any("version:" in t for t in base_tags)
assert any("POD_NAME:" in t for t in base_tags)
assert any("HOSTNAME:" in t for t in base_tags)
# Test with custom env vars
test_env = {
@ -705,12 +705,12 @@ def test_get_datadog_tags():
# Test with empty request_tags
standard_logging_obj["request_tags"] = []
tags_empty_request = get_datadog_tags(standard_logging_obj)
assert "request_tag:" not in tags_empty_request
assert not any(t.startswith("request_tag:") for t in tags_empty_request)
# Test with None request_tags
standard_logging_obj["request_tags"] = None
tags_none_request = get_datadog_tags(standard_logging_obj)
assert "request_tag:" not in tags_none_request
assert not any(t.startswith("request_tag:") for t in tags_none_request)
@pytest.mark.asyncio

2
tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py
View File

@ -44,7 +44,7 @@ class TestDatadogTagsRegression:
assert "env:test-env" in tags_legacy
assert "service:test-service" in tags_legacy
# Verify NO team tag (should not invent one)
assert "team:" not in tags_legacy
assert not any(t.startswith("team:") for t in tags_legacy)
# Case 2: New feature (team info provided)
payload_with_team = StandardLoggingPayload(

138
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View File

@ -1666,3 +1666,141 @@ async def test_oauth_authorize_prefers_request_scope_over_server_config():
redirect_url = response.headers["location"]
assert "scope=custom_scope1+custom_scope2" in redirect_url or "scope=custom_scope1%20custom_scope2" in redirect_url
assert "default_scope" not in redirect_url
@pytest.mark.asyncio
async def test_token_endpoint_refresh_token_grant():
"""Test that token endpoint supports refresh_token grant type."""
try:
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
token_endpoint,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
# Clear registry
global_mcp_server_manager.registry.clear()
# Create mock OAuth2 server
oauth2_server = MCPServer(
server_id="google_mcp",
name="google_mcp",
server_name="google_mcp",
alias="google_mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="test_client_id",
client_secret="test_secret",
authorization_url="https://accounts.google.com/o/oauth2/v2/auth",
token_url="https://oauth2.googleapis.com/token",
scopes=["openid", "email"],
)
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://proxy.litellm.example/"
mock_request.headers = {}
# Mock httpx client response with new tokens
mock_response = MagicMock()
mock_response.json.return_value = {
"access_token": "new_access_token",
"token_type": "Bearer",
"expires_in": 3599,
"refresh_token": "new_refresh_token",
}
mock_response.raise_for_status = MagicMock()
mock_async_client = MagicMock()
mock_async_client.post = AsyncMock(return_value=mock_response)
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = mock_async_client
response = await token_endpoint(
request=mock_request,
grant_type="refresh_token",
code=None,
redirect_uri=None,
client_id="test_client_id",
mcp_server_name="google_mcp",
client_secret="test_secret",
refresh_token="rt-test",
scope="openid email",
)
# Verify the POST was called with refresh_token grant data
mock_async_client.post.assert_called_once()
call_args = mock_async_client.post.call_args
assert call_args[1]["data"]["grant_type"] == "refresh_token"
assert call_args[1]["data"]["refresh_token"] == "rt-test"
assert call_args[1]["data"]["client_id"] == "test_client_id"
assert call_args[1]["data"]["client_secret"] == "test_secret"
assert call_args[1]["data"]["scope"] == "openid email"
# Verify response contains the new tokens
import json
token_data = json.loads(response.body)
assert token_data["access_token"] == "new_access_token"
assert token_data["refresh_token"] == "new_refresh_token"
@pytest.mark.asyncio
async def test_token_endpoint_authorization_code_missing_code():
"""Test that authorization_code grant rejects missing code param."""
try:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
exchange_token_with_server,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
global_mcp_server_manager.registry.clear()
server = MCPServer(
server_id="test_server",
name="test_server",
server_name="test_server",
alias="test_server",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
token_url="https://example.com/token",
)
global_mcp_server_manager.registry[server.server_id] = server
mock_request = MagicMock()
mock_request.base_url = "https://proxy.example/"
mock_request.headers = {}
with pytest.raises(HTTPException) as exc_info:
await exchange_token_with_server(
request=mock_request,
mcp_server=server,
grant_type="authorization_code",
code=None,
redirect_uri="https://example.com/cb",
client_id="cid",
client_secret=None,
code_verifier=None,
)
assert exc_info.value.status_code == 400
assert "code is required" in str(exc_info.value.detail)

88
tests/test_litellm/proxy/db/test_prisma_client.py
View File

@ -1,7 +1,8 @@
import json
import os
import signal
import sys
from unittest.mock import AsyncMock, Mock, patch
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
from fastapi.testclient import TestClient
@ -14,6 +15,14 @@ sys.path.insert(
from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema
@pytest.fixture(autouse=True)
def mock_prisma_binary():
"""Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests."""
mock_module = MagicMock()
with patch.dict(sys.modules, {"prisma": mock_module}):
yield mock_module
def test_should_update_prisma_schema(monkeypatch):
# CASE 1: Environment variable behavior
# When DISABLE_SCHEMA_UPDATE is not set -> should update
@ -73,4 +82,79 @@ async def test_recreate_prisma_client_successful_disconnect():
# Verify that the new client replaced the original
assert wrapper._original_prisma != mock_prisma
assert hasattr(wrapper._original_prisma, 'connect')
assert hasattr(wrapper._original_prisma, 'connect')
@pytest.mark.asyncio
async def test_recreate_prisma_client_kills_old_engine_on_disconnect_failure(
mock_prisma_binary,
):
"""When disconnect() fails, recreate_prisma_client must SIGTERM/SIGKILL the old engine PID."""
mock_prisma = AsyncMock()
mock_prisma.disconnect.side_effect = Exception("engine hung")
# Simulate engine subprocess with a known PID
mock_engine = MagicMock()
mock_engine.process.pid = 12345
mock_prisma._engine = mock_engine
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
# Configure the mock Prisma constructor
mock_new_prisma = AsyncMock()
mock_prisma_binary.Prisma.return_value = mock_new_prisma
with (
patch("os.kill") as mock_kill,
patch("asyncio.sleep", new_callable=AsyncMock),
):
await wrapper.recreate_prisma_client("postgresql://new")
# Verify old engine was killed
mock_kill.assert_any_call(12345, signal.SIGTERM)
# Verify new client was created and connected
mock_new_prisma.connect.assert_awaited_once()
@pytest.mark.asyncio
async def test_recreate_prisma_client_skips_kill_on_successful_disconnect(
mock_prisma_binary,
):
"""When disconnect() succeeds, no kill should be attempted."""
mock_prisma = AsyncMock()
mock_prisma.disconnect.return_value = None
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
mock_new_prisma = AsyncMock()
mock_prisma_binary.Prisma.return_value = mock_new_prisma
with patch("os.kill") as mock_kill:
await wrapper.recreate_prisma_client("postgresql://new")
mock_kill.assert_not_called()
mock_new_prisma.connect.assert_awaited_once()
@pytest.mark.asyncio
async def test_recreate_prisma_client_handles_missing_engine_pid(
mock_prisma_binary,
):
"""When engine PID is unavailable (no _engine attr), kill is skipped gracefully."""
mock_prisma = AsyncMock()
mock_prisma.disconnect.side_effect = Exception("engine hung")
mock_prisma._engine = None # No engine subprocess
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
mock_new_prisma = AsyncMock()
mock_prisma_binary.Prisma.return_value = mock_new_prisma
with (
patch("os.kill") as mock_kill,
patch("asyncio.sleep", new_callable=AsyncMock),
):
await wrapper.recreate_prisma_client("postgresql://new")
mock_kill.assert_not_called() # PID was 0, kill skipped
mock_new_prisma.connect.assert_awaited_once()

35
tests/test_litellm/proxy/db/test_prisma_self_heal.py
View File

@ -1,5 +1,6 @@
import asyncio
import os
import signal
import sys
import time
from unittest.mock import AsyncMock, MagicMock, patch
@ -279,3 +280,37 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
await client.stop_db_health_watchdog_task()
assert client._db_health_watchdog_task is None
assert dummy_task.cancelled() is True
@pytest.mark.asyncio
async def test_lightweight_reconnect_kills_engine_on_disconnect_failure(mock_proxy_logging):
"""Lightweight reconnect must kill the old engine PID when disconnect() fails."""
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client.db.disconnect = AsyncMock(side_effect=Exception("disconnect failed"))
client.db.connect = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
with (
patch.object(client, "_get_engine_pid", return_value=9999),
patch("os.kill") as mock_kill,
patch("asyncio.sleep", new_callable=AsyncMock),
):
await client._run_reconnect_cycle(timeout_seconds=5.0)
mock_kill.assert_any_call(9999, signal.SIGTERM)
client.db.connect.assert_awaited_once()
client.db.query_raw.assert_awaited_once_with("SELECT 1")
@pytest.mark.asyncio
async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(mock_proxy_logging):
"""Lightweight reconnect must NOT kill when disconnect() succeeds."""
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
with patch("os.kill") as mock_kill:
await client._run_reconnect_cycle(timeout_seconds=5.0)
mock_kill.assert_not_called()

52
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View File

@ -1519,6 +1519,8 @@ class TestTemporaryMCPSessionEndpoints:
client_id="client",
client_secret="secret",
code_verifier="verifier",
refresh_token=None,
scope=None,
)
assert result is exchange_response
@ -1532,6 +1534,56 @@ class TestTemporaryMCPSessionEndpoints:
client_id="client",
client_secret="secret",
code_verifier="verifier",
refresh_token=None,
scope=None,
)
@pytest.mark.asyncio
async def test_mcp_token_proxies_refresh_token_grant(self):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
mcp_token,
)
request = MagicMock()
server = generate_mock_mcp_server_config_record(server_id="server-1")
exchange_response = {"access_token": "new-token", "refresh_token": "new-rt"}
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404",
return_value=server,
) as get_server,
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server",
AsyncMock(return_value=exchange_response),
) as exchange_mock,
):
result = await mcp_token(
request=request,
server_id="server-1",
grant_type="refresh_token",
code=None,
redirect_uri=None,
client_id="client",
client_secret="secret",
code_verifier=None,
refresh_token="rt-123",
scope=None,
)
assert result is exchange_response
get_server.assert_called_once_with("server-1")
exchange_mock.assert_awaited_once_with(
request=request,
mcp_server=server,
grant_type="refresh_token",
code=None,
redirect_uri=None,
client_id="client",
client_secret="secret",
code_verifier=None,
refresh_token="rt-123",
scope=None,
)
@pytest.mark.asyncio