Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix-config-update-targeted-upserts

This commit is contained in:
Yuneng Jiang
2026-04-29 18:26:19 -07:00
134 changed files with 7968 additions and 1253 deletions
@@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead(
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive(
async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
engine_client,
) -> None:
"""_run_reconnect_cycle uses disconnect/connect when engine is alive."""
engine_client._engine_pid = 1234
"""Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1.
with patch.object(engine_client, "_is_engine_alive", return_value=True):
The old "lightweight" path called `disconnect()` + `connect()`, which
blocks the event loop on the sync `process.wait()` inside aclose().
The fix routes both engine-alive and engine-dead paths through
`recreate_prisma_client`, which non-blockingly kills the old engine.
"""
engine_client._engine_pid = 1234
engine_client._start_engine_watcher = AsyncMock()
with (
patch.object(engine_client, "_is_engine_alive", return_value=True),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.connect.assert_awaited_once()
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test"
)
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown(
async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown(
engine_client,
) -> None:
"""_run_reconnect_cycle uses lightweight path when engine PID is not tracked."""
"""When the engine PID is not tracked, direct reconnect still runs."""
engine_client._engine_pid = 0
engine_client._start_engine_watcher = AsyncMock()
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.connect.assert_awaited_once()
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test"
)
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
@@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client):
@pytest.mark.asyncio
async def test_escalation_after_consecutive_lightweight_failures(engine_client):
"""After N consecutive lightweight reconnect failures, _engine_confirmed_dead
async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client):
"""After N consecutive direct reconnect failures, _engine_confirmed_dead
is set to True so _run_reconnect_cycle takes the heavy reconnect path."""
engine_client._reconnect_escalation_threshold = 3
engine_client._consecutive_reconnect_failures = 0
engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test
engine_client._start_engine_watcher = AsyncMock(return_value=None)
# Make lightweight reconnect fail every time
engine_client.db.disconnect = AsyncMock(return_value=None)
engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed"))
# Make direct reconnect fail every time
engine_client.db.recreate_prisma_client = AsyncMock(
side_effect=Exception("recreate failed")
)
# Run 3 failed reconnect attempts
for i in range(3):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is False
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
for _ in range(3):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is False
assert engine_client._consecutive_reconnect_failures == 3
# Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle
# Next attempt should escalate to the heavy path (recreate_prisma_client still
# the call, but via the _engine_confirmed_dead branch that also re-arms the watcher).
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
engine_client._start_engine_watcher = AsyncMock(return_value=None)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test_escalation", timeout_seconds=5.0
)
# Heavy reconnect should have been attempted (recreate_prisma_client called)
engine_client.db.recreate_prisma_client.assert_awaited_once()
@@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client):
"""A successful reconnect resets _consecutive_reconnect_failures to 0."""
engine_client._consecutive_reconnect_failures = 2
engine_client._db_reconnect_cooldown_seconds = 0
engine_client._start_engine_watcher = AsyncMock()
# Make reconnect succeed
engine_client.db.disconnect = AsyncMock(return_value=None)
engine_client.db.connect = AsyncMock(return_value=None)
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is True
assert engine_client._consecutive_reconnect_failures == 0
@@ -314,11 +314,11 @@ def test_update_litellm_params_for_health_check():
# Issue #15807: Fixes health checks sending "region/model" as model ID to AWS
model_info = {}
litellm_params = {
"model": "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "anthropic.claude-3-7-sonnet-20250219-v1:0"
assert updated_params["model"] == "anthropic.claude-sonnet-4-5-20250929-v1:0"
# Test with Bedrock cross-region inference profile - should preserve the inference profile prefix
# AWS requires inference profile IDs like "us.anthropic.claude..." for cross-region routing
@@ -366,3 +366,40 @@ def test_generic_api_compatible_callbacks_json_unknown_callback():
# Should return the string unchanged
assert result == "unknown_callback", "Unknown callback should be returned as-is"
assert isinstance(result, str), "Unknown callback should remain a string"
@pytest.mark.asyncio
async def test_generic_api_callback_settings_retry_config():
"""
Test that generic_api callback_settings are passed to GenericAPILogger.
"""
from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
from litellm.litellm_core_utils.logging_callback_manager import (
_generic_api_logger_cache,
)
callback_name = "test_generic_api_retry_config"
_generic_api_logger_cache.pop(callback_name, None)
litellm.callback_settings[callback_name] = {
"callback_type": "generic_api",
"endpoint": "https://example.com/api/logs",
"headers": {"Content-Type": "application/json"},
"max_retries": 2,
"retry_delay": 0.5,
"timeout": 3,
}
try:
result = LoggingCallbackManager._add_custom_callback_generic_api_str(
callback_name
)
assert isinstance(result, GenericAPILogger)
assert result.endpoint == "https://example.com/api/logs"
assert result.headers == {"Content-Type": "application/json"}
assert result.max_retries == 2
assert result.retry_delay == 0.5
assert result.timeout == 3
finally:
litellm.callback_settings.pop(callback_name, None)
_generic_api_logger_cache.pop(callback_name, None)
+3 -3
View File
@@ -2309,11 +2309,11 @@ def test_get_provider_audio_transcription_config():
@pytest.mark.parametrize(
"model, expected_bool",
[
("anthropic.claude-3-7-sonnet-20250219-v1:0", True),
("us.anthropic.claude-3-7-sonnet-20250219-v1:0", True),
("anthropic.claude-sonnet-4-5-20250929-v1:0", True),
("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True),
],
)
def test_claude_3_7_sonnet_supports_pdf_input(model, expected_bool):
def test_claude_sonnet_4_5_supports_pdf_input(model, expected_bool):
from litellm.utils import supports_pdf_input
assert supports_pdf_input(model) == expected_bool
@@ -134,7 +134,7 @@ class TestBedrockAnthropicPromptCachingRegression:
if "converse" in model_prefix:
config = AmazonConverseConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -162,7 +162,7 @@ class TestBedrockAnthropicPromptCachingRegression:
else:
config = AmazonAnthropicClaudeConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -227,7 +227,7 @@ class TestBedrockAnthropicPromptCachingRegression:
if "converse" in model_prefix:
config = AmazonConverseConfig()
result = config._transform_request_helper(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
system_content_blocks=[],
optional_params={},
messages=messages,
@@ -236,7 +236,7 @@ class TestBedrockAnthropicPromptCachingRegression:
else:
config = AmazonAnthropicClaudeConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -498,7 +498,7 @@ class TestBedrockAnthropicCombinedRegressions:
if "converse" in model_prefix:
config = AmazonConverseConfig()
result = config._transform_request_helper(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
system_content_blocks=[],
optional_params={},
messages=messages,
@@ -518,7 +518,7 @@ class TestBedrockAnthropicCombinedRegressions:
else:
config = AmazonAnthropicClaudeConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -1323,7 +1323,7 @@ def test_base_aws_llm_get_credentials():
def test_bedrock_completion_test_2():
litellm.set_verbose = True
data = {
"model": "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
"messages": [
{
"role": "system",
@@ -1630,7 +1630,7 @@ def test_bedrock_completion_test_4(modify_params):
litellm.modify_params = modify_params
data = {
"model": "anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "anthropic.claude-sonnet-4-5-20250929-v1:0",
"messages": [
{
"role": "user",
@@ -2115,7 +2115,7 @@ class TestBedrockConverseAnthropicUnitTests(BaseAnthropicChatTest):
def get_base_completion_call_args_with_thinking(self) -> dict:
return {
"model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"thinking": {"type": "enabled", "budget_tokens": 16000},
}
@@ -2828,7 +2828,7 @@ async def test_bedrock_thinking_in_assistant_message(sync_mode):
client = AsyncHTTPHandler()
params = {
"model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"messages": [
{
"role": "assistant",
@@ -2887,7 +2887,7 @@ async def test_bedrock_stream_thinking_content_openwebui():
```
"""
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[{"role": "user", "content": "Hello who is this?"}],
stream=True,
max_tokens=1080,
@@ -580,7 +580,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers():
def test_litellm_gateway_from_sdk_with_thinking_param():
try:
response = litellm.completion(
model="litellm_proxy/anthropic.claude-3-7-sonnet-20250219-v1:0",
model="litellm_proxy/anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[{"role": "user", "content": "Hello world"}],
api_base="http://0.0.0.0:4000",
api_key="sk-PIp1h0RekR",
@@ -1828,7 +1828,7 @@ def test_azure_response_format_param():
"model, provider",
[
("claude-3-7-sonnet-20240620-v1:0", "anthropic"),
("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"),
("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"),
("invoke/anthropic.claude-3-7-sonnet-20240620-v1:0", "bedrock"),
("claude-3-7-sonnet@20250219", "vertex_ai"),
],
@@ -3493,8 +3493,14 @@ def test_litellm_api_base(monkeypatch, provider, route):
def test_gemini_tool_calling_working_demo():
load_vertex_ai_credentials()
litellm._turn_on_debug()
"""
Regression test: tool params with anyOf containing a `{"type": "array"}`
branch (no items field at all) must synthesize items before the request
is sent to Vertex (Vertex rejects array types missing items).
"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
args = {
"messages": [
{
@@ -3564,13 +3570,75 @@ def test_gemini_tool_calling_working_demo():
],
"vertex_location": "global",
}
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
print(response)
client = HTTPHandler()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello!"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15,
},
}
with (
patch.object(client, "post", return_value=mock_response) as mock_post,
patch.object(
VertexBase,
"_ensure_access_token",
return_value=("fake-token", "fake-project"),
),
):
completion(
model="vertex_ai/gemini-3-flash-preview",
client=client,
**args,
)
sent_body = mock_post.call_args.kwargs.get(
"json"
) or mock_post.call_args.kwargs.get("data")
assert sent_body is not None, "expected request body to be sent"
if isinstance(sent_body, str):
sent_body = json.loads(sent_body)
function_decl = sent_body["tools"][0]["function_declarations"][0]
callbacks_schema = function_decl["parameters"]["properties"]["config"][
"properties"
]["callbacks"]
array_branches = [
branch
for branch in callbacks_schema["anyOf"]
if branch.get("type", "").lower() == "array"
]
assert array_branches, "expected an array branch in callbacks anyOf"
for branch in array_branches:
assert "items" in branch and branch["items"], (
f"array branch in callbacks.anyOf must include non-empty items "
f"(Vertex rejects array types missing items). Got: {branch}"
)
def test_gemini_tool_calling_not_working():
load_vertex_ai_credentials()
litellm._turn_on_debug()
"""
Regression test: tool params with anyOf containing both an empty-items
array branch and a null branch must serialize with items present on the
array branch (Vertex rejects array types missing `items`).
"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
args = {
"messages": [
{
@@ -3637,8 +3705,64 @@ def test_gemini_tool_calling_not_working():
],
"vertex_location": "global",
}
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
print(response)
client = HTTPHandler()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello!"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15,
},
}
with (
patch.object(client, "post", return_value=mock_response) as mock_post,
patch.object(
VertexBase,
"_ensure_access_token",
return_value=("fake-token", "fake-project"),
),
):
completion(
model="vertex_ai/gemini-3-flash-preview",
client=client,
**args,
)
sent_body = mock_post.call_args.kwargs.get(
"json"
) or mock_post.call_args.kwargs.get("data")
assert sent_body is not None, "expected request body to be sent"
if isinstance(sent_body, str):
sent_body = json.loads(sent_body)
function_decl = sent_body["tools"][0]["function_declarations"][0]
callbacks_schema = function_decl["parameters"]["properties"]["config"][
"properties"
]["callbacks"]
array_branches = [
branch
for branch in callbacks_schema["anyOf"]
if branch.get("type", "").lower() == "array"
]
assert array_branches, "expected an array branch in callbacks anyOf"
for branch in array_branches:
assert "items" in branch and branch["items"], (
f"array branch in callbacks.anyOf must include non-empty items "
f"(Vertex rejects array types missing items). Got: {branch}"
)
def test_vertex_ai_llama_tool_calling():
+1 -1
View File
@@ -159,7 +159,7 @@ def test_aaparallel_function_call(model):
"model",
[
"anthropic/claude-4-sonnet-20250514",
"bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
],
)
@pytest.mark.flaky(retries=3, delay=1)
+3 -4
View File
@@ -30,10 +30,9 @@ def test_model_alias_map(caplog):
)
print(response.model)
captured_logs = [rec.levelname for rec in caplog.records]
for log in captured_logs:
assert "ERROR" not in log
for rec in caplog.records:
if rec.levelname == "ERROR" and rec.name.startswith("LiteLLM"):
pytest.fail(f"Unexpected litellm ERROR log: {rec.getMessage()}")
assert "llama-3.1-8b-instant" in response.model
except litellm.ServiceUnavailableError:
@@ -354,6 +354,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=litellm_params_dict,
extra_body=None,
)
)
captured_request_body["url"] = url
@@ -8,6 +8,7 @@ sys.path.insert(0, os.path.abspath("../.."))
import asyncio
import litellm
import gzip
import httpx
import json
import logging
import time
@@ -470,3 +471,96 @@ async def test_generic_api_callback_invalid_log_format():
endpoint=test_endpoint,
log_format="invalid_format", # type: ignore # Intentionally invalid for testing
)
@pytest.mark.asyncio
async def test_generic_api_callback_retries_timeout_then_succeeds():
"""
Test that GenericAPILogger retries LiteLLM timeout errors when configured.
"""
test_endpoint = "https://example.com/api/logs"
generic_logger = GenericAPILogger(
endpoint=test_endpoint,
max_retries=1,
retry_delay=0,
timeout=0.2,
)
mock_post = AsyncMock()
mock_post.side_effect = [
litellm.Timeout(
message="Connection timed out",
model="default-model-name",
llm_provider="litellm-httpx-handler",
),
type("Response", (), {"status_code": 200})(),
]
generic_logger.async_httpx_client.post = mock_post
generic_logger.log_queue = [{"event": "timeout-retry"}]
await generic_logger.async_send_batch()
assert mock_post.call_count == 2
first_call = mock_post.call_args_list[0][1]
assert first_call["url"] == test_endpoint
assert first_call["timeout"] == 0.2
assert json.loads(first_call["data"]) == [{"event": "timeout-retry"}]
@pytest.mark.asyncio
async def test_generic_api_callback_retries_5xx_then_succeeds():
"""
Test that GenericAPILogger retries transient HTTP 5xx errors when configured.
"""
test_endpoint = "https://example.com/api/logs"
generic_logger = GenericAPILogger(
endpoint=test_endpoint,
max_retries=1,
retry_delay=0,
)
request = httpx.Request("POST", test_endpoint)
response = httpx.Response(status_code=503, request=request)
mock_post = AsyncMock()
mock_post.side_effect = [
httpx.HTTPStatusError(
"Server error",
request=request,
response=response,
),
type("Response", (), {"status_code": 200})(),
]
generic_logger.async_httpx_client.post = mock_post
generic_logger.log_queue = [{"event": "5xx-retry"}]
await generic_logger.async_send_batch()
assert mock_post.call_count == 2
@pytest.mark.asyncio
async def test_generic_api_callback_does_not_retry_4xx():
"""
Test that GenericAPILogger does not retry non-transient HTTP 4xx errors.
"""
test_endpoint = "https://example.com/api/logs"
generic_logger = GenericAPILogger(
endpoint=test_endpoint,
max_retries=2,
retry_delay=0,
)
request = httpx.Request("POST", test_endpoint)
response = httpx.Response(status_code=401, request=request)
mock_post = AsyncMock()
mock_post.side_effect = httpx.HTTPStatusError(
"Unauthorized",
request=request,
response=response,
)
generic_logger.async_httpx_client.post = mock_post
generic_logger.log_queue = [{"event": "4xx-no-retry"}]
await generic_logger.async_send_batch()
mock_post.assert_called_once()
@@ -96,8 +96,8 @@ class BaseAnthropicMessagesPromptCachingTest(ABC):
Returns the model string to use for tests.
Examples:
- "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0"
- "bedrock/invoke/anthropic.claude-3-7-sonnet-20250219-v1:0"
- "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0"
- "bedrock/invoke/anthropic.claude-sonnet-4-5-20250929-v1:0"
"""
pass
@@ -31,7 +31,7 @@ class TestBedrockConversePromptCaching(BaseAnthropicMessagesPromptCachingTest):
"""
def get_model(self) -> str:
return "bedrock/converse/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
return "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest):
@@ -43,4 +43,4 @@ class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest):
"""
def get_model(self) -> str:
return "bedrock/invoke/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
return "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
@@ -52,3 +52,183 @@ async def test_process_async_embedding_cached_response():
print(f"response: {response}")
assert len(response.data) == 1
@pytest.mark.asyncio
async def test_embedding_cache_preserves_prompt_tokens_details():
"""Test that prompt_tokens_details (including image_count) survives a full cache hit."""
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
cached_result = [
{
"embedding": [-0.025, -0.019],
"index": 0,
"object": "embedding",
"model": "amazon.titan-embed-image-v1",
"prompt_tokens_details": {"image_count": 1},
}
]
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
final_embedding_cached_response=None,
cached_result=cached_result,
kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"},
logging_obj=mock_logging_obj,
start_time=datetime.now(),
model="amazon.titan-embed-image-v1",
)
assert cache_hit
assert response.usage is not None
assert response.usage.prompt_tokens_details is not None
assert response.usage.prompt_tokens_details.image_count == 1
@pytest.mark.asyncio
async def test_embedding_cache_backward_compat_no_prompt_tokens_details():
"""Test that old cached items without prompt_tokens_details still work."""
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
# Old-format cached item — no prompt_tokens_details field
cached_result = [
{
"embedding": [-0.025, -0.019],
"index": 0,
"object": "embedding",
"model": "text-embedding-ada-002",
}
]
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
final_embedding_cached_response=None,
cached_result=cached_result,
kwargs={"model": "text-embedding-ada-002", "input": "test"},
logging_obj=mock_logging_obj,
start_time=datetime.now(),
model="text-embedding-ada-002",
)
assert cache_hit
assert response.usage is not None
assert response.usage.prompt_tokens_details is None
@pytest.mark.asyncio
async def test_embedding_cache_aggregates_multiple_image_counts():
"""Test that image_count is summed correctly across multiple cached items."""
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
cached_result = [
{
"embedding": [-0.025, -0.019],
"index": 0,
"object": "embedding",
"model": "amazon.titan-embed-image-v1",
"prompt_tokens_details": {"image_count": 1},
},
{
"embedding": [0.031, 0.042],
"index": 1,
"object": "embedding",
"model": "amazon.titan-embed-image-v1",
"prompt_tokens_details": {"image_count": 1},
},
]
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
final_embedding_cached_response=None,
cached_result=cached_result,
kwargs={
"model": "amazon.titan-embed-image-v1",
"input": ["img1", "img2"],
},
logging_obj=mock_logging_obj,
start_time=datetime.now(),
model="amazon.titan-embed-image-v1",
)
assert cache_hit
assert response.usage.prompt_tokens_details is not None
assert response.usage.prompt_tokens_details.image_count == 2
def test_combine_usage_merges_prompt_tokens_details():
"""Test that combine_usage merges prompt_tokens_details from both Usage objects."""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
usage1 = Usage(
prompt_tokens=10,
completion_tokens=0,
total_tokens=10,
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
)
usage2 = Usage(
prompt_tokens=20,
completion_tokens=0,
total_tokens=20,
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2),
)
combined = llm_caching_handler.combine_usage(usage1, usage2)
assert combined.prompt_tokens == 30
assert combined.total_tokens == 30
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.image_count == 3
def test_combine_usage_handles_none_details():
"""Test that combine_usage works when one or both sides have null prompt_tokens_details."""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
# Both null
usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10)
usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20)
combined = llm_caching_handler.combine_usage(usage_a, usage_b)
assert combined.prompt_tokens_details is None
# Only first has details
usage_c = Usage(
prompt_tokens=10,
completion_tokens=0,
total_tokens=10,
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
)
combined = llm_caching_handler.combine_usage(usage_c, usage_b)
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.image_count == 1
# Only second has details
combined = llm_caching_handler.combine_usage(usage_a, usage_c)
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.image_count == 1
File diff suppressed because one or more lines are too long
@@ -1 +1 @@
{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}
{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}
@@ -220,7 +220,7 @@ async def test_anthropic_cache_control_hook_negative_indices():
with patch.object(client, "post", return_value=mock_response) as mock_post:
# Test with multiple messages and negative indices
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "system",
@@ -352,7 +352,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds_logging():
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
cache_control_injection_points=[
{"location": "message", "index": 10}
@@ -420,7 +420,7 @@ async def test_anthropic_cache_control_hook_negative_out_of_bounds_logging():
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
cache_control_injection_points=[
{
@@ -486,7 +486,7 @@ async def test_anthropic_cache_control_hook_multiple_user_messages():
with patch.object(client, "post", return_value=mock_response) as mock_post:
# Test with multiple user messages and negative indices
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "user",
@@ -586,7 +586,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds(bad_index):
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
cache_control_injection_points=[
{"location": "message", "index": bad_index}
@@ -651,7 +651,7 @@ async def test_anthropic_cache_control_hook_single_message(message_list):
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=message_list,
cache_control_injection_points=[{"location": "message", "index": -1}],
client=client,
@@ -691,7 +691,7 @@ async def test_anthropic_cache_control_hook_empty_message_list():
match="bedrock requires at least one non-system message",
):
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[],
cache_control_injection_points=[
{"location": "message", "index": -1}
@@ -742,7 +742,7 @@ async def test_anthropic_cache_control_hook_no_op():
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
# No cache_control_injection_points parameter
client=client,
@@ -799,7 +799,7 @@ async def test_anthropic_cache_control_hook_multiple_content_items_last_only():
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "user",
@@ -874,7 +874,7 @@ async def test_anthropic_cache_control_hook_document_analysis_multiple_pages():
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "user",
@@ -1057,7 +1057,7 @@ async def test_anthropic_cache_control_hook_string_negative_index():
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{"role": "user", "content": "First message"},
{"role": "assistant", "content": "First response"},
@@ -262,7 +262,7 @@ class TestOpenTelemetryProviderInitialization(unittest.TestCase):
class TestOpenTelemetry(unittest.TestCase):
POLL_INTERVAL = 0.05
POLL_TIMEOUT = 2.0
MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
HERE = os.path.dirname(__file__)
@patch.dict(os.environ, {}, clear=True)
@@ -369,7 +369,7 @@ def test_generic_cost_per_token_gpt55():
def test_generic_cost_per_token_gpt55_pro():
"""gpt-5.5-pro: responses-only model — $60/1M input, $360/1M output, $6/1M cached input."""
"""gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input."""
model = "gpt-5.5-pro"
custom_llm_provider = "openai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@@ -378,18 +378,18 @@ def test_generic_cost_per_token_gpt55_pro():
model_cost_map = litellm.model_cost[model]
# Sanity-check the map values match OpenAI's published pricing.
assert model_cost_map["input_cost_per_token"] == 6e-5
assert model_cost_map["output_cost_per_token"] == 3.6e-4
assert model_cost_map["cache_read_input_token_cost"] == 6e-6
assert model_cost_map["input_cost_per_token"] == 3e-5
assert model_cost_map["output_cost_per_token"] == 1.8e-4
assert model_cost_map["cache_read_input_token_cost"] == 3e-6
assert model_cost_map["litellm_provider"] == "openai"
# gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint).
assert model_cost_map["mode"] == "responses"
assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"]
assert "/v1/responses" in model_cost_map["supported_endpoints"]
# Inherits GPT-5.4-pro's long-context window + tiered pricing (scaled 2x).
# Inherits GPT-5.4-pro's long-context window + tiered pricing.
assert model_cost_map["max_input_tokens"] == 1050000
assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1.2e-4
assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 5.4e-4
assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 6e-5
assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 2.7e-4
prompt_tokens = 1000
completion_tokens = 500
@@ -454,8 +454,8 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(
[
("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7),
("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7),
("azure/gpt-5.5-pro", "responses", 6e-5, 3.6e-4, 6e-6),
("azure/gpt-5.5-pro-2026-04-23", "responses", 6e-5, 3.6e-4, 6e-6),
("azure/gpt-5.5-pro", "responses", 3e-5, 1.8e-4, 3e-6),
("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6),
],
)
def test_azure_gpt55_entries_present_with_correct_pricing(
@@ -464,7 +464,7 @@ def test_azure_gpt55_entries_present_with_correct_pricing(
"""Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure.
Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page
on 2026-04-24): $5/$30 input/output per 1M for chat, $60/$360 for pro.
on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro.
Cache discount is 10% of input.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@@ -77,7 +77,7 @@ async def test_anthropic_bedrock_thinking_blocks_with_none_content():
# test _bedrock_converse_messages_pt_async
result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
llm_provider="bedrock",
)
@@ -0,0 +1,138 @@
"""
Regression tests for the parsed-URL hostname match used to identify a
caller-supplied ``api_base`` as a known openai-compatible provider.
The previous shape (``if endpoint in api_base:``) used unanchored
substring search, which let a caller pass
``https://attacker.com/api.groq.com/openai/v1`` and have the proxy
return ``GROQ_API_KEY`` as the dynamic credential exfiltrating the
server's real provider key to an attacker-controlled host on the
outbound request.
"""
import os
import sys
from unittest.mock import patch
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.litellm_core_utils.get_llm_provider_logic import (
_endpoint_matches_api_base,
get_llm_provider,
)
class TestEndpointMatchesApiBase:
"""Direct unit tests on the parsed-URL matcher."""
@pytest.mark.parametrize(
"endpoint, api_base",
[
# Bare hostname endpoint, exact host match.
("api.perplexity.ai", "https://api.perplexity.ai/v1"),
# Endpoint includes a path; api_base path starts with it.
("api.groq.com/openai/v1", "https://api.groq.com/openai/v1"),
# Endpoint with full URL scheme.
("https://api.cerebras.ai/v1", "https://api.cerebras.ai/v1/chat"),
# Trailing-slash on registered endpoint must not break match.
("https://llm.chutes.ai/v1/", "https://llm.chutes.ai/v1/chat"),
# Case-insensitive on hostname.
("api.groq.com/openai/v1", "https://API.GROQ.COM/openai/v1"),
],
)
def test_legitimate_provider_urls_match(self, endpoint, api_base):
assert _endpoint_matches_api_base(endpoint, api_base) is True
@pytest.mark.parametrize(
"endpoint, api_base",
[
# Attacker host, registered endpoint smuggled into path.
(
"api.groq.com/openai/v1",
"https://attacker.com/api.groq.com/openai/v1",
),
# Attacker host, registered endpoint smuggled into a path segment.
(
"api.groq.com/openai/v1",
"https://attacker.com/foo/api.groq.com/openai/v1",
),
# Lookalike host that contains the registered host as a suffix label.
(
"api.groq.com/openai/v1",
"https://api.groq.com.attacker.com/openai/v1",
),
# Lookalike host with the registered host as a prefix.
(
"api.groq.com/openai/v1",
"https://api.groq.com.evil.example/openai/v1",
),
# Right host, wrong path — endpoint requires ``/openai/v1`` prefix.
("api.groq.com/openai/v1", "https://api.groq.com/v1"),
# Path-segment lookalike: ``/openai/v10`` must not match ``/openai/v1``.
("api.groq.com/openai/v1", "https://api.groq.com/openai/v10"),
# Userinfo / @-injection trick — the ``hostname`` after ``@`` is
# what httpx connects to.
(
"api.groq.com/openai/v1",
"https://api.groq.com@attacker.com/openai/v1",
),
],
)
def test_attacker_smuggling_does_not_match(self, endpoint, api_base):
assert _endpoint_matches_api_base(endpoint, api_base) is False
class TestGetLlmProviderRejectsAttackerSmuggledApiBase:
"""
End-to-end: ``get_llm_provider`` must NOT return the server's stored
secret (e.g. ``GROQ_API_KEY``) for an api_base whose hostname is
attacker-controlled, even when the registered endpoint string appears
elsewhere in the URL.
"""
def test_attacker_host_does_not_yield_groq_secret(self):
# The function may either fall through (different provider) or
# raise BadRequestError because the model can't be identified.
# The invariant under test is that ``GROQ_API_KEY`` is never
# looked up against an attacker-controlled hostname.
import litellm
with patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
return_value="server-real-groq-key",
) as mocked_secret:
try:
_, _, dynamic_api_key, _ = get_llm_provider(
model="some-model",
api_base="https://attacker.com/api.groq.com/openai/v1",
)
# If it returned, the dynamic key must not be the secret.
assert dynamic_api_key != "server-real-groq-key"
except litellm.exceptions.BadRequestError:
# Acceptable outcome: provider unidentifiable, no secret
# was returned.
pass
# Regardless of return / raise, the secret must never have been
# read against this attacker-controlled api_base.
groq_lookups = [
call
for call in mocked_secret.call_args_list
if call.args and call.args[0] == "GROQ_API_KEY"
]
assert groq_lookups == []
def test_legitimate_groq_api_base_still_resolves(self):
with patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
return_value="server-real-groq-key",
):
_, provider, dynamic_api_key, _ = get_llm_provider(
model="some-model",
api_base="https://api.groq.com/openai/v1",
)
assert provider == "groq"
assert dynamic_api_key == "server-real-groq-key"
@@ -2337,6 +2337,104 @@ def test_merge_hidden_params_from_response_into_metadata_populates_metadata():
assert meta["hidden_params"]["model_id"] == "mid-test"
def test_merge_hidden_params_from_response_into_metadata_backfills_response_cost():
"""Streaming metadata should include the already-calculated response cost."""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="merge-hp-cost-test",
function_id="merge-hp-cost-fn",
)
logging_obj.model_call_details = {
"litellm_params": {"metadata": {}},
"response_cost": 0.002,
}
class _Resp:
_hidden_params = {"response_cost": None, "model_id": "mid-test"}
response = _Resp()
logging_obj._merge_hidden_params_from_response_into_metadata(response)
meta = logging_obj.model_call_details["litellm_params"]["metadata"]
assert meta["hidden_params"]["response_cost"] == 0.002
assert meta["hidden_params"]["model_id"] == "mid-test"
assert response._hidden_params["response_cost"] is None
def test_standard_logging_hidden_params_backfills_response_cost_without_mutating_response():
"""Streaming standard logging payload should expose the calculated response cost."""
from datetime import datetime
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import Usage
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="standard-hp-cost-test",
function_id="standard-hp-cost-fn",
)
logging_obj.model_call_details = {
"litellm_params": {"metadata": {}, "proxy_server_request": {}},
"litellm_call_id": "standard-hp-cost-test",
"call_type": "acompletion",
"stream": True,
"model": "gpt-4o-mini",
"custom_llm_provider": "openai",
"optional_params": {"stream": True},
"response_cost": 0.002,
}
response = ModelResponse(
id="standard-hp-cost-response",
model="gpt-4o-mini",
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
response._hidden_params = {"response_cost": None, "model_id": "mid-test"}
payload = logging_obj._build_standard_logging_payload(
response, datetime.now(), datetime.now()
)
assert payload is not None
assert payload["hidden_params"]["response_cost"] == 0.002
assert response._hidden_params["response_cost"] is None
def test_merge_hidden_params_from_response_into_metadata_preserves_response_cost():
"""Do not overwrite provider-supplied response cost when it already exists."""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="merge-hp-preserve-cost-test",
function_id="merge-hp-preserve-cost-fn",
)
logging_obj.model_call_details = {
"litellm_params": {"metadata": {}},
"response_cost": 0.002,
}
class _Resp:
_hidden_params = {"response_cost": 0.001, "model_id": "mid-test"}
logging_obj._merge_hidden_params_from_response_into_metadata(_Resp())
meta = logging_obj.model_call_details["litellm_params"]["metadata"]
assert meta["hidden_params"]["response_cost"] == 0.001
assert meta["hidden_params"]["model_id"] == "mid-test"
def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty():
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -2436,3 +2534,141 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
assert payload is not None
assert payload["litellm_call_id"] == call_id
def _make_dict_logging_obj():
"""Build a Logging instance configured for a non-streaming dict result."""
obj = LitellmLogging(
model="claude-haiku-4-5@20251001",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="acompletion",
litellm_call_id="test-call-id",
start_time=time.time(),
function_id="test-fn",
)
obj.model_call_details = {
"model": "claude-haiku-4-5@20251001",
"custom_llm_provider": "vertex_ai",
"litellm_params": {"metadata": {}},
"response_cost": None,
}
return obj
def test_success_handler_computes_cost_for_dict_response():
"""Non-streaming dict responses run through the cost calculator."""
logging_obj = _make_dict_logging_obj()
expected_cost = 0.42
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=expected_cost,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=False,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result={"id": "msg_1"},
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_called_once()
assert logging_obj.model_call_details["response_cost"] == expected_cost
def test_success_handler_preserves_precomputed_cost_for_dict_response():
"""Precomputed response_cost on model_call_details must not be overwritten."""
logging_obj = _make_dict_logging_obj()
precomputed_cost = 1.23
logging_obj.model_call_details["response_cost"] = precomputed_cost
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=9.99,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": precomputed_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=False,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result={"id": "msg_2"},
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_not_called()
assert logging_obj.model_call_details["response_cost"] == precomputed_cost
def test_success_handler_unified_helper_runs_for_typed_results():
"""Recognized typed responses still flow through the unified helper."""
logging_obj = _make_dict_logging_obj()
expected_cost = 0.10
typed_result = MagicMock()
typed_result._hidden_params = {}
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=expected_cost,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=True,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result=typed_result,
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_called_once()
assert logging_obj.model_call_details["response_cost"] == expected_cost
@@ -279,7 +279,7 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto():
}
optional_params = config.map_openai_params(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
non_default_params=non_default_params,
optional_params={},
drop_params=False,
@@ -2797,7 +2797,7 @@ def test_thinking_with_max_completion_tokens():
result = config.map_openai_params(
non_default_params=non_default_params_with_max_completion,
optional_params=optional_params,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
@@ -2819,7 +2819,7 @@ def test_thinking_with_max_completion_tokens():
result = config.map_openai_params(
non_default_params=non_default_params_with_max_tokens,
optional_params=optional_params,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
@@ -2842,7 +2842,7 @@ def test_thinking_with_max_completion_tokens():
result = config.map_openai_params(
non_default_params=non_default_params_without_max,
optional_params=optional_params,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
@@ -3617,7 +3617,7 @@ class TestBedrockMinThinkingBudgetTokens:
"""Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024)."""
def _map_params(
self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0"
self, thinking_value, model="anthropic.claude-sonnet-4-5-20250929-v1:0"
):
"""Helper to call map_openai_params with the given thinking value."""
config = AmazonConverseConfig()
@@ -3651,7 +3651,7 @@ class TestBedrockMinThinkingBudgetTokens:
result = config.map_openai_params(
non_default_params={},
optional_params={},
model="anthropic.claude-3-7-sonnet-20250219-v1:0",
model="anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
assert "thinking" not in result or result.get("thinking") is None
@@ -4146,3 +4146,39 @@ def test_transform_response_finish_reason_stop_when_json_mode_filters_all_tools(
# finish_reason must be "stop", not "tool_calls"
assert result.choices[0].finish_reason == "stop"
def test_transform_response_does_not_leak_body_on_parse_failure():
from litellm.llms.bedrock.common_utils import BedrockError
leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}}
class MockResponse:
def json(self):
return leaky_body
@property
def text(self):
return json.dumps(leaky_body)
with patch(
"litellm.llms.bedrock.chat.converse_transformation.ConverseResponseBlock",
side_effect=KeyError("missing required field"),
):
with pytest.raises(BedrockError) as exc_info:
AmazonConverseConfig()._transform_response(
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
response=MockResponse(),
model_response=ModelResponse(),
stream=False,
logging_obj=None,
optional_params={},
api_key=None,
data=None,
messages=[],
encoding=None,
)
msg = str(exc_info.value)
assert "secret content" not in msg
assert "Error converting to valid response block" in msg
@@ -21,7 +21,112 @@ def test_transform_search_request():
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body=None,
)
assert url.endswith("/kb123/retrieve")
assert body["retrievalQuery"].get("text") == "hello"
def test_transform_search_request_uses_only_retrieval_config_from_extra_body():
config = BedrockVectorStoreConfig()
mock_log = MagicMock()
mock_log.model_call_details = {}
url, body = config.transform_search_vector_store_request(
vector_store_id="kb123",
query="hello",
vector_store_search_optional_params={},
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body={
"retrievalConfiguration": {
"vectorSearchConfiguration": {
"overrideSearchType": "HYBRID",
"numberOfResults": 8,
}
},
"unrelatedField": {"should_not": "be_forwarded"},
},
)
assert url.endswith("/kb123/retrieve")
assert body["retrievalQuery"].get("text") == "hello"
assert (
body["retrievalConfiguration"]["vectorSearchConfiguration"][
"overrideSearchType"
]
== "HYBRID"
)
assert "unrelatedField" not in body
def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results():
config = BedrockVectorStoreConfig()
mock_log = MagicMock()
mock_log.model_call_details = {}
extra_body = {
"retrievalConfiguration": {
"vectorSearchConfiguration": {
"overrideSearchType": "HYBRID",
"numberOfResults": 8,
}
}
}
_, body = config.transform_search_vector_store_request(
vector_store_id="kb123",
query="hello",
vector_store_search_optional_params={"max_num_results": 10},
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body=extra_body,
)
assert (
body["retrievalConfiguration"]["vectorSearchConfiguration"]["numberOfResults"]
== 10
)
assert (
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"][
"numberOfResults"
]
== 8
)
def test_transform_search_request_overrides_filter_without_mutating_extra_body():
config = BedrockVectorStoreConfig()
mock_log = MagicMock()
mock_log.model_call_details = {}
extra_body = {
"retrievalConfiguration": {
"vectorSearchConfiguration": {
"filter": {"equals": {"key": "tenant", "value": "a"}}
}
}
}
new_filter = {"equals": {"key": "tenant", "value": "b"}}
_, body = config.transform_search_vector_store_request(
vector_store_id="kb123",
query="hello",
vector_store_search_optional_params={"filters": new_filter},
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body=extra_body,
)
assert (
body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"]
== new_filter
)
assert (
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"][
"equals"
]["value"]
== "a"
)
@@ -0,0 +1,161 @@
import socket
from unittest.mock import MagicMock, patch
def _invoke_connector_factory(http_handler_module):
"""
Drive the lambda factory installed on the transport so TCPConnector is
actually constructed. _create_aiohttp_transport returns a transport whose
_client_factory is the lambda that builds (TCPConnector ClientSession);
invoking it directly avoids relying on _get_valid_client_session's internal
branching to trigger connector construction.
"""
transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
shared_session=None
)
transport._client_factory()
return transport
def test_socket_factory_omitted_when_disabled(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", False)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
connector_mock = MagicMock(name="connector")
session_mock = MagicMock(name="session")
with patch.object(
http_handler_module, "TCPConnector", return_value=connector_mock
) as mock_tcp_connector:
with patch.object(
http_handler_module, "ClientSession", return_value=session_mock
):
_invoke_connector_factory(http_handler_module)
assert mock_tcp_connector.call_count >= 1
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
def test_socket_factory_attached_when_enabled(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
connector_mock = MagicMock(name="connector")
session_mock = MagicMock(name="session")
with patch.object(
http_handler_module, "TCPConnector", return_value=connector_mock
) as mock_tcp_connector:
with patch.object(
http_handler_module, "ClientSession", return_value=session_mock
):
_invoke_connector_factory(http_handler_module)
assert mock_tcp_connector.call_count >= 1
factory = mock_tcp_connector.call_args.kwargs.get("socket_factory")
assert callable(factory)
def test_socket_factory_skipped_on_old_aiohttp(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", False)
connector_mock = MagicMock(name="connector")
session_mock = MagicMock(name="session")
with patch.object(
http_handler_module, "TCPConnector", return_value=connector_mock
) as mock_tcp_connector:
with patch.object(
http_handler_module, "ClientSession", return_value=session_mock
):
_invoke_connector_factory(http_handler_module)
assert mock_tcp_connector.call_count >= 1
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
def test_socket_factory_sets_keepalive_options(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 45)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPINTVL", 15)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPCNT", 4)
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
assert factory is not None
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
fake_sock = MagicMock(spec=socket.socket)
with patch("socket.socket", return_value=fake_sock) as sock_ctor:
returned = factory(addr_info)
sock_ctor.assert_called_once_with(
family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP
)
assert returned is fake_sock
fake_sock.setblocking.assert_called_once_with(False)
setsockopt_calls = {
(call.args[0], call.args[1]): call.args[2]
for call in fake_sock.setsockopt.call_args_list
}
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
if hasattr(socket, "TCP_KEEPIDLE"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)] == 45
elif hasattr(socket, "TCP_KEEPALIVE"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)] == 45
if hasattr(socket, "TCP_KEEPINTVL"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)] == 15
if hasattr(socket, "TCP_KEEPCNT"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)] == 4
def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch):
"""
Cover the macOS/Darwin branch: when TCP_KEEPIDLE is missing but TCP_KEEPALIVE
is present, the factory should fall back to TCP_KEEPALIVE for the idle timer.
Linux CI runners always have TCP_KEEPIDLE, so we patch socket itself to
simulate the BSD-derived environment.
"""
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 60)
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
assert factory is not None
fake_socket_module = MagicMock(spec=[])
fake_socket_module.SOL_SOCKET = socket.SOL_SOCKET
fake_socket_module.SO_KEEPALIVE = socket.SO_KEEPALIVE
fake_socket_module.IPPROTO_TCP = socket.IPPROTO_TCP
fake_socket_module.TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10)
fake_sock = MagicMock(spec=socket.socket)
fake_socket_module.socket = MagicMock(return_value=fake_sock)
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
with patch.object(http_handler_module, "socket", fake_socket_module):
factory(addr_info)
setsockopt_calls = {
(call.args[0], call.args[1]): call.args[2]
for call in fake_sock.setsockopt.call_args_list
}
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
assert (
setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60
)
assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls
@@ -2,12 +2,16 @@ import os
import sys
from unittest.mock import AsyncMock, Mock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import (
BaseLLMHTTPHandler,
_google_genai_streaming_hidden_params,
)
from litellm.types.router import GenericLiteLLMParams
@@ -320,3 +324,29 @@ async def test_async_anthropic_messages_handler_header_priority():
assert captured_headers["X-Forwarded-Only"] == "keep"
assert captured_headers["X-Extra-Only"] == "also-keep"
assert captured_headers["X-Provider-Only"] == "keep-this-too"
def test_google_genai_streaming_hidden_params_model_info_and_router_fallback():
logging_obj = Mock()
logging_obj.get_router_model_id = Mock(return_value="router-model-id")
from_model_info = _google_genai_streaming_hidden_params(
api_base="https://generativelanguage.googleapis.com/v1beta",
litellm_params=GenericLiteLLMParams(model_info={"id": "info-id"}),
logging_obj=logging_obj,
response_headers=httpx.Headers({"x-ratelimit-remaining": "10"}),
)
assert from_model_info["model_id"] == "info-id"
assert (
from_model_info["api_base"]
== "https://generativelanguage.googleapis.com/v1beta"
)
assert isinstance(from_model_info["additional_headers"], dict)
from_router = _google_genai_streaming_hidden_params(
api_base="https://x",
litellm_params=GenericLiteLLMParams(),
logging_obj=logging_obj,
response_headers=httpx.Headers({}),
)
assert from_router["model_id"] == "router-model-id"
@@ -59,6 +59,7 @@ class TestS3VectorsVectorStoreConfig:
api_base="https://s3vectors.us-west-2.api.aws",
litellm_logging_obj=mock_logging_obj,
litellm_params={},
extra_body=None,
)
def test_transform_search_response(self):
@@ -4261,3 +4261,32 @@ def test_sync_streaming_uses_custom_client():
# Verify that gemini_client is in the partial's keywords
assert "gemini_client" in partial_make_sync_call.keywords
assert partial_make_sync_call.keywords["gemini_client"] is mock_client
def test_transform_response_does_not_leak_body_on_parse_failure():
leaky_body = {"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]}
raw_response = MagicMock()
raw_response.json.return_value = leaky_body
raw_response.text = json.dumps(leaky_body)
raw_response.headers = {}
with patch(
"litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.GenerateContentResponseBody",
side_effect=KeyError("missing required field"),
):
with pytest.raises(VertexAIError) as exc_info:
VertexGeminiConfig().transform_response(
model="gemini-pro",
raw_response=raw_response,
model_response=ModelResponse(),
logging_obj=MagicMock(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
msg = str(exc_info.value)
assert "secret content" not in msg
assert "Error converting to valid response block" in msg
@@ -225,7 +225,11 @@ def test_build_vertex_schema():
"metadata": {"type": "object"},
"callbacks": {
"anyOf": [
{"type": "array", "nullable": True},
{
"type": "array",
"items": {"type": "object"},
"nullable": True,
},
{"type": "object", "nullable": True},
]
},
@@ -288,6 +292,43 @@ def test_process_items_basic():
process_items(schema)
assert schema["properties"]["nested"]["items"] == {"type": "object"}
# Vertex rejects array types missing `items` entirely (not just empty).
# Synthesize {"type": "object"} so the request validates.
schema = {"type": "array"}
process_items(schema)
assert schema["items"] == {"type": "object"}
def test_build_vertex_schema_array_branch_missing_items_in_anyof():
"""
Regression: an `anyOf` branch with `{"type": "array"}` (no items) must
end up with synthesized `items: {"type": "object"}` after the schema
transform Vertex returns INVALID_ARGUMENT otherwise.
"""
from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
parameters = {
"properties": {
"callbacks": {
"anyOf": [
{"type": "array"},
{"type": "object"},
{"type": "null"},
]
}
},
"type": "object",
}
result = _build_vertex_schema(parameters)
callbacks_anyof = result["properties"]["callbacks"]["anyOf"]
array_branches = [b for b in callbacks_anyof if b.get("type") == "array"]
assert array_branches, "expected an array branch to remain after transform"
for branch in array_branches:
assert branch.get("items") == {
"type": "object"
}, f"array branch must have items synthesized; got {branch}"
def test_vertex_ai_complex_response_schema():
import json
@@ -311,3 +311,29 @@ def test_transform_anthropic_messages_request_removes_scope_from_cache_control()
# scope removed from message content
assert "scope" not in result["messages"][0]["content"][0]["cache_control"]
assert result["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral"
def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance():
"""
Regression test: repeated provider config lookups for the same Vertex Claude model
should return the same config instance (which preserves auth cache state).
"""
import litellm
from litellm.utils import ProviderConfigManager
ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear()
try:
first_config = ProviderConfigManager.get_provider_anthropic_messages_config(
model="claude-opus-4-6",
provider=litellm.LlmProviders.VERTEX_AI,
)
second_config = ProviderConfigManager.get_provider_anthropic_messages_config(
model="claude-opus-4-6",
provider=litellm.LlmProviders.VERTEX_AI,
)
assert isinstance(first_config, VertexAIPartnerModelsAnthropicMessagesConfig)
assert isinstance(second_config, VertexAIPartnerModelsAnthropicMessagesConfig)
assert first_config is second_config
finally:
ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear()
@@ -728,6 +728,136 @@ class TestMCPServerManager:
]
assert scopes == ["read", "write"]
@pytest.mark.asyncio
async def test_descovery_metadata_probes_well_known_when_server_does_not_challenge(
self,
):
manager = MCPServerManager()
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock()
mock_client = MagicMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_metadata = MCPOAuthMetadata(
scopes=None,
authorization_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize",
token_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/token",
registration_url=None,
)
with (
patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
return_value=mock_client,
),
patch.object(
manager,
"_attempt_well_known_discovery",
AsyncMock(
return_value=(
["https://login.microsoftonline.com/test-tenant-id/v2.0"],
["api://some-scope/.default"],
)
),
) as mock_well_known,
patch.object(
manager,
"_fetch_authorization_server_metadata",
AsyncMock(return_value=mock_metadata),
) as mock_fetch_auth,
):
result = await manager._descovery_metadata("http://localhost:8001/mcp")
mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp")
mock_fetch_auth.assert_awaited_once_with(
["https://login.microsoftonline.com/test-tenant-id/v2.0"]
)
assert result is mock_metadata
assert result.scopes == ["api://some-scope/.default"]
@pytest.mark.asyncio
async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path(
self,
):
manager = MCPServerManager()
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
def build_response(url: str):
mock_response = MagicMock()
if url == f"{issuer}/.well-known/openid-configuration":
mock_response.json.return_value = {
"authorization_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize",
"token_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token",
"scopes_supported": ["api://some-scope/.default"],
}
mock_response.raise_for_status = MagicMock()
else:
request = httpx.Request("GET", url)
response_obj = httpx.Response(status_code=404, request=request)
mock_response.raise_for_status = MagicMock(
side_effect=httpx.HTTPStatusError(
"not found", request=request, response=response_obj
)
)
return mock_response
mock_client = MagicMock()
mock_client.get = AsyncMock(side_effect=build_response)
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
return_value=mock_client,
):
result = await manager._fetch_single_authorization_server_metadata(issuer)
assert result is not None
assert (
result.authorization_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
)
assert (
result.token_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
)
assert result.scopes == ["api://some-scope/.default"]
@pytest.mark.asyncio
async def test_fetch_single_authorization_server_metadata_derives_azure_metadata(
self,
):
manager = MCPServerManager()
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
request = httpx.Request("GET", issuer)
response_obj = httpx.Response(status_code=404, request=request)
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock(
side_effect=httpx.HTTPStatusError(
"not found", request=request, response=response_obj
)
)
mock_client = MagicMock()
mock_client.get = AsyncMock(return_value=mock_response)
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
return_value=mock_client,
):
result = await manager._fetch_single_authorization_server_metadata(issuer)
assert result is not None
assert (
result.authorization_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
)
assert (
result.token_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
)
@pytest.mark.asyncio
async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self):
manager = MCPServerManager()
@@ -0,0 +1,302 @@
"""
Tests for the short-ID MCP tool prefix (LITELLM_USE_SHORT_MCP_TOOL_PREFIX).
The short-prefix mode swaps the historical alias/server_name prefix on
tool names for a deterministic three-character base62 ID derived from the
server's ``server_id``. This keeps tool names well below the 60-char
upper bound enforced by some model APIs while remaining stable across
processes/restarts and tolerant of mixed-version clients.
"""
from typing import List
import pytest
from mcp.types import Tool as MCPTool
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
from litellm.proxy._experimental.mcp_server.utils import (
SHORT_MCP_TOOL_PREFIX_LENGTH,
add_server_prefix_to_name,
compute_short_server_prefix,
get_server_prefix,
is_short_mcp_tool_prefix_enabled,
iter_known_server_prefixes,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
def _make_server(
*,
server_id: str = "abcdef-1234",
server_name: str = "github_onprem",
alias: str = "github_onprem",
) -> MCPServer:
return MCPServer(
server_id=server_id,
name=alias or server_name,
alias=alias,
server_name=server_name,
transport="http",
)
@pytest.fixture(autouse=True)
def _reset_env(monkeypatch):
monkeypatch.delenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", raising=False)
yield
# ---------------------------------------------------------------------------
# Pure helpers
# ---------------------------------------------------------------------------
class TestShortPrefixHelpers:
def test_short_prefix_is_three_base62_chars(self):
prefix = compute_short_server_prefix("any-server-id")
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
assert prefix.isalnum() and prefix.isascii()
def test_short_prefix_first_char_is_alphabetic(self):
"""The first char must be [A-Za-z] so the prefix is a valid identifier
on every model API (some providers historically required the first
character of a function name to be alphabetic)."""
# Sweep many server_ids and rehash attempts to give us coverage of
# every position the high-order bits can land on.
for i in range(200):
for attempt in range(4):
prefix = compute_short_server_prefix(f"server-{i}", attempt=attempt)
assert prefix[0].isalpha(), (
f"prefix {prefix!r} for server-{i} (attempt={attempt}) "
f"starts with a non-alphabetic character"
)
def test_short_prefix_is_deterministic(self):
assert compute_short_server_prefix("abc") == compute_short_server_prefix("abc")
assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd")
def test_short_prefix_requires_server_id(self):
with pytest.raises(ValueError):
compute_short_server_prefix("")
def test_flag_defaults_to_false(self):
assert is_short_mcp_tool_prefix_enabled() is False
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "On"])
def test_flag_truthy_values(self, monkeypatch, value):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
assert is_short_mcp_tool_prefix_enabled() is True
@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
def test_flag_falsey_values(self, monkeypatch, value):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
assert is_short_mcp_tool_prefix_enabled() is False
# ---------------------------------------------------------------------------
# get_server_prefix behaviour
# ---------------------------------------------------------------------------
class TestGetServerPrefix:
def test_default_mode_uses_alias(self):
server = _make_server(alias="github_onprem", server_name="github_onprem")
assert get_server_prefix(server) == "github_onprem"
def test_short_mode_uses_short_id(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = _make_server(server_id="abcdef-1234")
prefix = get_server_prefix(server)
assert prefix == compute_short_server_prefix("abcdef-1234")
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
def test_short_mode_falls_back_when_no_server_id(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
class _Bare:
alias = "fallback_alias"
server_name = None
server_id = None
assert get_server_prefix(_Bare()) == "fallback_alias"
# ---------------------------------------------------------------------------
# iter_known_server_prefixes — covers reverse-lookup tolerance
# ---------------------------------------------------------------------------
class TestIterKnownServerPrefixes:
def test_default_mode_includes_short_id_too(self):
server = _make_server()
prefixes = list(iter_known_server_prefixes(server))
# Contains the live prefix and every known form so that mixed-mode
# clients can be resolved.
assert "github_onprem" in prefixes
assert compute_short_server_prefix(server.server_id) in prefixes
def test_short_mode_still_yields_long_forms(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = _make_server()
prefixes = list(iter_known_server_prefixes(server))
assert "github_onprem" in prefixes
assert compute_short_server_prefix(server.server_id) in prefixes
# ---------------------------------------------------------------------------
# Manager-level behaviour: list + reverse-lookup
# ---------------------------------------------------------------------------
def _stub_tools() -> List[MCPTool]:
return [
MCPTool(name="get_repo", description="", inputSchema={"type": "object"}),
MCPTool(name="list_issues", description="", inputSchema={"type": "object"}),
]
class TestManagerShortPrefix:
def test_list_tools_uses_short_prefix_when_flag_on(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
manager = MCPServerManager()
server = _make_server()
out = manager._create_prefixed_tools(_stub_tools(), server)
short = compute_short_server_prefix(server.server_id)
assert {t.name for t in out} == {f"{short}-get_repo", f"{short}-list_issues"}
def test_call_tool_lookup_resolves_short_prefix(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
manager = MCPServerManager()
server = _make_server()
manager.registry[server.server_id] = server
manager._create_prefixed_tools(_stub_tools(), server)
short = compute_short_server_prefix(server.server_id)
resolved = manager._get_mcp_server_from_tool_name(f"{short}-get_repo")
assert resolved is server
def test_call_tool_lookup_resolves_long_prefix_in_short_mode(self, monkeypatch):
"""Old clients that cached the long-prefix name must still route."""
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
manager = MCPServerManager()
server = _make_server()
manager.registry[server.server_id] = server
manager._create_prefixed_tools(_stub_tools(), server)
resolved = manager._get_mcp_server_from_tool_name("github_onprem-get_repo")
assert resolved is server
def test_default_mode_unchanged(self):
manager = MCPServerManager()
server = _make_server()
out = manager._create_prefixed_tools(_stub_tools(), server)
assert {t.name for t in out} == {
"github_onprem-get_repo",
"github_onprem-list_issues",
}
assert (
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is None
) # registry empty
manager.registry[server.server_id] = server
assert (
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is server
)
def test_total_tool_name_length_short_enough(self, monkeypatch):
"""The short prefix keeps tool names under the 60-char limit even
when the upstream tool name is itself reasonably long."""
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
long_server_name = "a" * 50
server = _make_server(
server_id="server-id-1",
server_name=long_server_name,
alias=long_server_name,
)
prefix = get_server_prefix(server)
full = add_server_prefix_to_name("get_repo", prefix)
assert len(full) < 60
# ---------------------------------------------------------------------------
# Collision-resolution at registration time
# ---------------------------------------------------------------------------
class TestShortPrefixCollisionResolution:
"""``_assign_unique_short_prefix`` must rehash on collision.
The dedup path is exercised by forcing two distinct ``server_id``
values to both hash to the same natural prefix via a monkeypatched
``compute_short_server_prefix``.
"""
def test_no_op_when_flag_off(self):
manager = MCPServerManager()
server = _make_server(server_id="abc")
manager._assign_unique_short_prefix(server)
assert server.short_prefix is None
def test_assigns_natural_hash_when_no_collision(self, monkeypatch):
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
manager = MCPServerManager()
server = _make_server(server_id="abc")
manager._assign_unique_short_prefix(server)
assert server.short_prefix == mcp_utils.compute_short_server_prefix("abc")
def test_rehashes_when_natural_hash_collides(self, monkeypatch):
"""Two server_ids that natural-hash to the same prefix get
deterministic, distinct short prefixes."""
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
# Force every attempt=0 hash to "AAA" and attempt=1 to "AAB".
# That way the second server registered must rehash to "AAB".
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
def _fake_hash(server_id: str, attempt: int = 0) -> str:
return "AAA" if attempt == 0 else f"AA{chr(ord('A') + attempt)}"
monkeypatch.setattr(mcp_utils, "compute_short_server_prefix", _fake_hash)
# Also patch the symbol that the manager imported at module load.
from litellm.proxy._experimental.mcp_server import (
mcp_server_manager as mgr_module,
)
monkeypatch.setattr(mgr_module, "compute_short_server_prefix", _fake_hash)
manager = MCPServerManager()
first = _make_server(server_id="server-1", alias="srv1")
second = _make_server(server_id="server-2", alias="srv2")
# Pretend both are already in the registry so dedup sees both.
manager.registry[first.server_id] = first
manager._assign_unique_short_prefix(first)
manager.registry[second.server_id] = second
manager._assign_unique_short_prefix(second)
assert first.short_prefix == "AAA"
assert second.short_prefix == "AAB"
assert first.short_prefix != second.short_prefix
def test_cached_prefix_is_reused(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
manager = MCPServerManager()
server = _make_server(server_id="abc")
server.short_prefix = "ZZZ" # pretend a previous registration set this
manager._assign_unique_short_prefix(server)
assert server.short_prefix == "ZZZ"
def test_get_server_prefix_prefers_cached(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = _make_server(server_id="abc")
server.short_prefix = "Q9q"
assert get_server_prefix(server) == "Q9q"
@@ -5,6 +5,8 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
from typing import Optional
from unittest.mock import MagicMock, patch
import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
_get_customer_id_from_standard_headers,
@@ -15,6 +17,7 @@ from litellm.proxy.auth.auth_utils import (
get_key_model_tpm_limit,
get_project_model_rpm_limit,
get_project_model_tpm_limit,
is_request_body_safe,
)
@@ -660,3 +663,304 @@ class TestCheckCompleteCredentials:
def test_returns_true_when_api_key_is_valid(self):
result = check_complete_credentials({"model": "gpt-4", "api_key": "sk-valid"})
assert result is True
class TestCheckCompleteCredentialsBlocksSSRF:
"""
Even with credentials supplied, ``api_base`` / ``base_url`` must not
point at private / internal / cloud-metadata addresses. Without this
the gate accepts ``api_key=anything`` plus a malicious target and the
proxy is used as an SSRF pivot.
The check only runs when ``litellm.user_url_validation`` is True, so
every test in this class flips the toggle. Tests stay mock-only no
real DNS is performed.
"""
@pytest.fixture(autouse=True)
def _enable_url_validation(self, monkeypatch):
import litellm
monkeypatch.setattr(litellm, "user_url_validation", True, raising=False)
@pytest.mark.parametrize(
"url_field",
["api_base", "base_url"],
)
@pytest.mark.parametrize(
"blocked_url",
[
"http://169.254.169.254/latest/meta-data/iam/security-credentials/",
"http://metadata.google.internal/computeMetadata/v1/",
"http://127.0.0.1:8080/admin",
"http://10.0.0.1/",
"http://192.168.1.1/",
],
)
def test_rejects_private_or_metadata_targets(self, url_field, blocked_url):
from litellm.litellm_core_utils.url_utils import SSRFError
with patch(
"litellm.proxy.auth.auth_utils.validate_url",
side_effect=SSRFError(f"blocked: {blocked_url}"),
):
with pytest.raises(ValueError) as exc_info:
check_complete_credentials(
{
"model": "gpt-4",
"api_key": "sk-some-clientside-key",
url_field: blocked_url,
}
)
assert url_field in str(exc_info.value)
assert "SSRF" in str(exc_info.value)
def test_allows_public_target_when_validate_url_passes(self):
# ``validate_url`` is mocked so no real DNS is performed.
with patch(
"litellm.proxy.auth.auth_utils.validate_url",
return_value=("https://api.openai.com/v1", "api.openai.com"),
):
result = check_complete_credentials(
{
"model": "gpt-4",
"api_key": "sk-some-clientside-key",
"api_base": "https://api.openai.com/v1",
}
)
assert result is True
def test_skips_url_validation_when_toggle_is_off(self, monkeypatch):
# Admins who disable ``user_url_validation`` (default) should not
# have requests rejected at the proxy boundary even if the URL
# would fail the SSRF guard.
import litellm
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
with patch(
"litellm.proxy.auth.auth_utils.validate_url",
) as mocked:
result = check_complete_credentials(
{
"model": "gpt-4",
"api_key": "sk-some-clientside-key",
"api_base": "http://127.0.0.1:8080/admin",
}
)
assert result is True
mocked.assert_not_called()
class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
"""
When the caller redirects ``api_base`` / ``base_url`` to their own
server, admin-set fields like ``OpenAI-Organization``, ``extra_body``,
AWS / Vertex / Azure tokens, and per-deployment ``api_version`` must
NOT flow through to that destination.
"""
def test_clears_admin_organization_and_extra_body_on_base_override(self):
from litellm.router_utils.clientside_credential_handler import (
get_dynamic_litellm_params,
)
admin_params = {
"model": "gpt-4",
"api_key": "sk-admin-key",
"api_base": "https://admin.upstream/v1",
"organization": "org-admin-corp",
"extra_body": {"x-admin-secret": "super-secret"},
"api_version": "2026-04-01",
}
out = get_dynamic_litellm_params(
litellm_params=dict(admin_params),
request_kwargs={
"api_key": "sk-attacker",
"api_base": "https://attacker.example",
},
)
assert out["api_base"] == "https://attacker.example"
assert out["api_key"] == "sk-attacker"
assert "organization" not in out
assert "extra_body" not in out
assert "api_version" not in out
def test_clears_aws_and_vertex_secrets_on_base_override(self):
from litellm.router_utils.clientside_credential_handler import (
get_dynamic_litellm_params,
)
admin_params = {
"model": "bedrock/claude-3",
"aws_access_key_id": "AKIA-EXAMPLE",
"aws_secret_access_key": "secret-example",
"aws_session_token": "session-example",
"vertex_credentials": '{"private_key":"-----BEGIN..."}',
"vertex_project": "admin-gcp-project",
}
out = get_dynamic_litellm_params(
litellm_params=dict(admin_params),
request_kwargs={"base_url": "https://attacker.example"},
)
assert "aws_access_key_id" not in out
assert "aws_secret_access_key" not in out
assert "aws_session_token" not in out
assert "vertex_credentials" not in out
assert "vertex_project" not in out
def test_caller_resupplied_value_overrides_admin_value_on_base_override(self):
# When the caller redirects ``api_base`` and *also* supplies their
# own value for one of the admin fields (e.g. ``organization``),
# the caller's value must win — never the admin's. The naive
# ``if field not in request_kwargs: pop`` shape lets a caller echo
# the field name with any value (or empty string) to keep the
# admin's value forwarded, which is the exfiltration vector this
# test guards against.
from litellm.router_utils.clientside_credential_handler import (
get_dynamic_litellm_params,
)
out = get_dynamic_litellm_params(
litellm_params={
"api_base": "https://admin.upstream/v1",
"organization": "org-admin",
"extra_body": {"admin": "value"},
},
request_kwargs={
"api_base": "https://attacker.example",
"organization": "org-attacker",
"extra_body": {"attacker": "value"},
},
)
assert out["organization"] == "org-attacker"
assert out["extra_body"] == {"attacker": "value"}
def test_field_echo_does_not_preserve_admin_value(self):
# Regression: a caller that echoes an admin-config field name with
# an *empty* value (or any value) must not be able to keep the
# admin's value in ``litellm_params``.
from litellm.router_utils.clientside_credential_handler import (
get_dynamic_litellm_params,
)
out = get_dynamic_litellm_params(
litellm_params={
"api_base": "https://admin.upstream/v1",
"organization": "org-admin-secret",
"extra_body": {"x-admin-only": "secret"},
},
request_kwargs={
"api_base": "https://attacker.example",
"organization": "",
"extra_body": "",
},
)
assert out["organization"] == ""
assert out["extra_body"] == ""
assert "org-admin-secret" not in str(out)
def test_no_clearing_when_only_api_key_overridden(self):
from litellm.router_utils.clientside_credential_handler import (
get_dynamic_litellm_params,
)
# Caller only overrides api_key (BYOK pattern); admin's organization /
# extra_body / region still apply because the destination is unchanged.
out = get_dynamic_litellm_params(
litellm_params={
"api_base": "https://admin.upstream/v1",
"organization": "org-admin",
"api_version": "2026-04-01",
},
request_kwargs={"api_key": "sk-byok"},
)
assert out["organization"] == "org-admin"
assert out["api_version"] == "2026-04-01"
assert out["api_base"] == "https://admin.upstream/v1"
class TestIsRequestBodySafeBlocksEndpointTargetingFields:
"""
``is_request_body_safe`` rejects request-body fields that retarget the
outbound request to a caller-controlled host. Beyond the original
``api_base`` / ``base_url``, the same protection must apply to:
* ``aws_bedrock_runtime_endpoint`` Bedrock endpoint redirect; an
attacker-controlled value coerces the proxy to authenticate against
their host with the admin's AWS creds.
* ``langsmith_base_url`` Langsmith callback host; attacker-controlled
values exfiltrate the entire request payload (incl. message content)
via the observability hook.
* ``langfuse_host`` same exfil vector via the Langfuse hook.
"""
@pytest.fixture(autouse=True)
def _disable_url_validation(self, monkeypatch):
# The new banned-params entries should be rejected even when
# ``user_url_validation`` is off — the gate isn't the URL guard,
# it's the banned-params list.
import litellm
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
@pytest.mark.parametrize(
"field",
[
"aws_bedrock_runtime_endpoint",
"langsmith_base_url",
"langfuse_host",
"posthog_host",
"braintrust_host",
"slack_webhook_url",
"s3_endpoint_url",
"sagemaker_base_url",
"deployment_url",
],
)
def test_endpoint_targeting_field_in_request_body_is_rejected(self, field):
with pytest.raises(ValueError) as exc:
is_request_body_safe(
request_body={"model": "gpt-4", field: "https://attacker.example"},
general_settings={},
llm_router=None,
model="gpt-4",
)
# The function lists the offending param name in the error.
assert field in str(exc.value)
@pytest.mark.parametrize(
"field",
["api_base", "base_url", "user_config", "langfuse_host", "slack_webhook_url"],
)
def test_api_key_does_not_bypass_blocklist(self, field):
# Regression: the historical ``check_complete_credentials`` clause
# made the entire blocklist a no-op for any caller that supplied
# a non-empty ``api_key``. That bypass turned every missing entry
# on the blocklist into an SSRF / credential-exfil hole. Verify
# that supplying an api_key (alongside the banned param) does NOT
# bypass the gate — it can only be opened by an admin opt-in.
with pytest.raises(ValueError) as exc:
is_request_body_safe(
request_body={
"model": "gpt-4",
"api_key": "sk-anything",
field: "https://attacker.example",
},
general_settings={},
llm_router=None,
model="gpt-4",
)
assert field in str(exc.value)
def test_admin_opt_in_proxy_wide_still_allows(self):
# ``general_settings.allow_client_side_credentials = True`` remains
# the documented proxy-wide BYOK opt-in.
assert (
is_request_body_safe(
request_body={"model": "gpt-4", "api_base": "https://my-byok.example"},
general_settings={"allow_client_side_credentials": True},
llm_router=None,
model="gpt-4",
)
is True
)
@@ -0,0 +1,348 @@
"""
Test expired UI session key cleanup manager functionality.
"""
import os
import sys
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException, status
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.constants import (
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
UI_SESSION_TOKEN_TEAM_ID,
)
from litellm.proxy._types import LiteLLM_VerificationToken
from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import (
ExpiredUISessionKeyCleanupManager,
)
class TestExpiredUISessionKeyCleanupManager:
"""Test the ExpiredUISessionKeyCleanupManager class functionality."""
@pytest.mark.asyncio
async def test_find_expired_ui_session_keys_filters_dashboard_team_and_expiry(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
now = datetime(2026, 4, 25, 12, 0, 0, tzinfo=timezone.utc)
mock_keys = [
LiteLLM_VerificationToken(
token="expired-dashboard-token",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=now - timedelta(seconds=1),
)
]
mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = (
mock_keys
)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.datetime"
) as mock_datetime:
mock_datetime.now.return_value = now
mock_datetime.side_effect = lambda *args, **kwargs: datetime(
*args, **kwargs
)
keys = await manager._find_expired_ui_session_keys()
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
where={
"team_id": UI_SESSION_TOKEN_TEAM_ID,
"expires": {"lt": now},
},
take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
)
assert keys == mock_keys
@pytest.mark.asyncio
async def test_cleanup_expired_keys_uses_existing_delete_path(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
expired_key = LiteLLM_VerificationToken(
token="expired-dashboard-token",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
)
manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key])
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
new_callable=AsyncMock,
) as mock_delete_verification_tokens:
mock_delete_verification_tokens.return_value = (
{"deleted_keys": ["expired-dashboard-token"], "failed_tokens": []},
[expired_key],
)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
new_callable=AsyncMock,
) as mock_key_deleted_hook:
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 1
mock_delete_verification_tokens.assert_called_once()
call_kwargs = mock_delete_verification_tokens.call_args.kwargs
assert call_kwargs["tokens"] == ["expired-dashboard-token"]
assert call_kwargs["user_api_key_cache"] == mock_cache
assert (
call_kwargs["litellm_changed_by"]
== LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
)
assert call_kwargs["user_api_key_dict"].user_id == "system"
mock_key_deleted_hook.assert_called_once()
hook_kwargs = mock_key_deleted_hook.call_args.kwargs
assert hook_kwargs["data"].keys == ["expired-dashboard-token"]
assert hook_kwargs["keys_being_deleted"] == [expired_key]
assert hook_kwargs["response"] == {
"deleted_keys": ["expired-dashboard-token"],
"failed_tokens": [],
}
assert (
hook_kwargs["litellm_changed_by"]
== LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
)
@pytest.mark.asyncio
async def test_cleanup_expired_keys_deletes_multiple_keys(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
expired_keys = [
LiteLLM_VerificationToken(
token="expired-dashboard-token-1",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
),
LiteLLM_VerificationToken(
token="expired-dashboard-token-2",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
),
]
tokens = [key.token for key in expired_keys]
manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
new_callable=AsyncMock,
) as mock_delete_verification_tokens:
mock_delete_verification_tokens.return_value = (
{"deleted_keys": tokens, "failed_tokens": []},
expired_keys,
)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
new_callable=AsyncMock,
) as mock_key_deleted_hook:
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 2
assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens
hook_kwargs = mock_key_deleted_hook.call_args.kwargs
assert hook_kwargs["data"].keys == tokens
assert hook_kwargs["keys_being_deleted"] == expired_keys
assert hook_kwargs["response"] == {"deleted_keys": tokens, "failed_tokens": []}
@pytest.mark.asyncio
async def test_cleanup_expired_keys_returns_successful_delete_count(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
expired_keys = [
LiteLLM_VerificationToken(
token="expired-dashboard-token-1",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
),
LiteLLM_VerificationToken(
token="expired-dashboard-token-2",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
),
]
tokens = [key.token for key in expired_keys]
manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
new_callable=AsyncMock,
) as mock_delete_verification_tokens:
mock_delete_verification_tokens.return_value = (
{
"deleted_keys": ["expired-dashboard-token-1"],
"failed_tokens": ["expired-dashboard-token-2"],
},
[expired_keys[0]],
)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
new_callable=AsyncMock,
):
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 1
assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens
@pytest.mark.asyncio
async def test_cleanup_expired_keys_counts_nested_delete_response(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
expired_keys = [
LiteLLM_VerificationToken(
token="expired-dashboard-token-1",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
),
LiteLLM_VerificationToken(
token="expired-dashboard-token-2",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
),
]
tokens = [key.token for key in expired_keys]
manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
new_callable=AsyncMock,
) as mock_delete_verification_tokens:
mock_delete_verification_tokens.return_value = (
{
"deleted_keys": {"deleted_keys": 2},
"failed_tokens": tokens,
},
expired_keys,
)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
new_callable=AsyncMock,
):
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 2
@pytest.mark.asyncio
async def test_cleanup_expired_keys_treats_missing_keys_as_noop(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
expired_key = LiteLLM_VerificationToken(
token="expired-dashboard-token",
team_id=UI_SESSION_TOKEN_TEAM_ID,
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
)
manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key])
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
new_callable=AsyncMock,
) as mock_delete_verification_tokens:
mock_delete_verification_tokens.side_effect = HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail={"error": "No keys found"},
)
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
new_callable=AsyncMock,
) as mock_key_deleted_hook:
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 0
mock_key_deleted_hook.assert_not_called()
@pytest.mark.asyncio
async def test_cleanup_expired_keys_noops_when_no_keys_found(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
manager._find_expired_ui_session_keys = AsyncMock(return_value=[])
with patch(
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
new_callable=AsyncMock,
) as mock_delete_verification_tokens:
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 0
mock_delete_verification_tokens.assert_not_called()
@pytest.mark.asyncio
async def test_cleanup_expired_keys_skips_when_lock_held(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
mock_pod_lock_manager = MagicMock()
mock_pod_lock_manager.redis_cache = MagicMock()
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
mock_pod_lock_manager.release_lock = AsyncMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
pod_lock_manager=mock_pod_lock_manager,
)
manager._find_expired_ui_session_keys = AsyncMock()
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 0
mock_pod_lock_manager.acquire_lock.assert_called_once_with(
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
)
manager._find_expired_ui_session_keys.assert_not_called()
mock_pod_lock_manager.release_lock.assert_not_called()
@pytest.mark.asyncio
async def test_cleanup_expired_keys_releases_acquired_lock(self):
mock_prisma_client = AsyncMock()
mock_cache = MagicMock()
mock_pod_lock_manager = MagicMock()
mock_pod_lock_manager.redis_cache = MagicMock()
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
mock_pod_lock_manager.release_lock = AsyncMock()
manager = ExpiredUISessionKeyCleanupManager(
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
pod_lock_manager=mock_pod_lock_manager,
)
manager._find_expired_ui_session_keys = AsyncMock(return_value=[])
deleted_count = await manager.cleanup_expired_keys()
assert deleted_count == 0
mock_pod_lock_manager.release_lock.assert_called_once_with(
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
)
@@ -0,0 +1,255 @@
"""
Unit tests for `call_with_db_reconnect_retry` the canonical "try DB read,
on transport error reconnect once and retry once" helper.
Covers the regression in issue #25143 where read paths (e.g.
`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in
LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient
`httpx.ReadError` flaps that used to self-heal in 1.82.6.
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from prisma.errors import UniqueViolationError
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
def _make_client(
*,
attempt_db_reconnect_return: bool = True,
has_attempt_db_reconnect: bool = True,
):
"""Build a minimal stand-in for PrismaClient that exposes only the surface
`call_with_db_reconnect_retry` actually pokes at."""
client = MagicMock()
if has_attempt_db_reconnect:
client.attempt_db_reconnect = AsyncMock(
return_value=attempt_db_reconnect_return
)
else:
# `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock
# auto-creates attributes, so we wipe it out via `spec`.
client = MagicMock(spec=[])
client._db_auth_reconnect_timeout_seconds = 2.0
client._db_auth_reconnect_lock_timeout_seconds = 0.1
return client
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_returns_value_on_first_success():
"""Happy path: factory succeeds first call, no reconnect attempted."""
client = _make_client()
async def _factory():
return {"id": 1}
result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path")
assert result == {"id": 1}
client.attempt_db_reconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_retries_after_transport_error():
"""Transport error on first call → reconnect → second call succeeds."""
client = _make_client(attempt_db_reconnect_return=True)
invocations = []
async def _factory():
invocations.append(None)
if len(invocations) == 1:
raise httpx.ReadError("transport blip")
return {"id": 1}
result = await call_with_db_reconnect_retry(
client, _factory, reason="prisma_get_generic_data_config_lookup_failure"
)
assert result == {"id": 1}
assert len(invocations) == 2
client.attempt_db_reconnect.assert_awaited_once()
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error():
"""Data-layer errors (e.g. UniqueViolationError) are NOT transport errors —
propagate immediately, do not reconnect."""
client = _make_client()
async def _factory():
raise UniqueViolationError(
data={"user_facing_error": {"meta": {}}},
message="Unique constraint failed",
)
with pytest.raises(UniqueViolationError):
await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test")
client.attempt_db_reconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails():
"""Transport error, but reconnect returns False → propagate the original
exception. Do not call factory a second time."""
client = _make_client(attempt_db_reconnect_return=False)
invocations = []
async def _factory():
invocations.append(None)
raise httpx.ReadError("transport blip")
with pytest.raises(httpx.ReadError):
await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails")
assert len(invocations) == 1
client.attempt_db_reconnect.assert_awaited_once()
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error():
"""Transport error, reconnect succeeds, retry also raises transport error →
propagate. At most one retry by construction (no infinite loop)."""
client = _make_client(attempt_db_reconnect_return=True)
invocations = []
async def _factory():
invocations.append(None)
raise httpx.ReadError("still failing")
with pytest.raises(httpx.ReadError):
await call_with_db_reconnect_retry(
client, _factory, reason="second_transport_error"
)
assert len(invocations) == 2
client.attempt_db_reconnect.assert_awaited_once()
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr():
"""Older PrismaClient stand-ins / partial mocks may not expose
`attempt_db_reconnect`. The helper must not crash just propagate the
original exception. Mirrors the `hasattr` guard from
`auth_checks._fetch_key_object_from_db_with_reconnect`."""
client = _make_client(has_attempt_db_reconnect=False)
async def _factory():
raise httpx.ReadError("transport blip")
with pytest.raises(httpx.ReadError):
await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr")
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro():
"""Guard against the obvious bug of awaiting the same coroutine twice
(`RuntimeError: cannot reuse already awaited coroutine`). The helper must
call the factory a fresh time on retry, not cache an awaitable."""
client = _make_client(attempt_db_reconnect_return=True)
factory_call_count = 0
async def _factory():
nonlocal factory_call_count
factory_call_count += 1
if factory_call_count == 1:
raise httpx.ReadError("transport blip")
return "ok"
result = await call_with_db_reconnect_retry(
client, _factory, reason="fresh_coro_on_retry"
)
assert result == "ok"
assert factory_call_count == 2
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_passes_explicit_timeouts():
"""Explicit timeout_seconds / lock_timeout_seconds override the auth
defaults read off the prisma_client object."""
client = _make_client(attempt_db_reconnect_return=True)
async def _factory():
if not hasattr(_factory, "_called"):
_factory._called = True # type: ignore[attr-defined]
raise httpx.ReadError("transport blip")
return "ok"
result = await call_with_db_reconnect_retry(
client,
_factory,
reason="explicit_timeouts",
timeout_seconds=5.5,
lock_timeout_seconds=0.25,
)
assert result == "ok"
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
assert call_kwargs["timeout_seconds"] == 5.5
assert call_kwargs["lock_timeout_seconds"] == 0.25
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset():
"""When timeouts are not provided, helper reads
`_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`
off the prisma_client (matching the auth path's existing convention)."""
client = _make_client(attempt_db_reconnect_return=True)
client._db_auth_reconnect_timeout_seconds = 3.0
client._db_auth_reconnect_lock_timeout_seconds = 0.5
async def _factory():
if not hasattr(_factory, "_called"):
_factory._called = True # type: ignore[attr-defined]
raise httpx.ReadError("transport blip")
return "ok"
await call_with_db_reconnect_retry(client, _factory, reason="defaults")
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
assert call_kwargs["timeout_seconds"] == 3.0
assert call_kwargs["lock_timeout_seconds"] == 0.5
@pytest.mark.asyncio
async def test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises():
"""If `attempt_db_reconnect` itself raises (lock cancellation, timer
error, unexpected internal failure), the helper must surface the
*original* transport error to telemetry not the reconnect exception.
Otherwise `failure_handler` / `db_exceptions` alerts log the wrong
error string and the actual DB transport problem becomes invisible.
The reconnect error is chained as the `__cause__` for debuggability."""
client = MagicMock()
reconnect_exc = RuntimeError("simulated reconnect lock cancellation")
client.attempt_db_reconnect = AsyncMock(side_effect=reconnect_exc)
client._db_auth_reconnect_timeout_seconds = 2.0
client._db_auth_reconnect_lock_timeout_seconds = 0.1
original_exc = httpx.ReadError("transport blip")
async def _factory():
raise original_exc
with pytest.raises(httpx.ReadError) as exc_info:
await call_with_db_reconnect_retry(
client, _factory, reason="reconnect_itself_raises"
)
assert exc_info.value is original_exc
assert exc_info.value.__cause__ is reconnect_exc
client.attempt_db_reconnect.assert_awaited_once()
@@ -5,6 +5,7 @@ import sys
import time
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
sys.path.insert(
@@ -34,18 +35,18 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging):
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.recreate_prisma_client = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
client._start_engine_watcher = AsyncMock()
result = await client.attempt_db_reconnect(
reason="unit_test_reconnect_success",
force=True,
)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await client.attempt_db_reconnect(
reason="unit_test_reconnect_success",
force=True,
)
assert result is True
client.db.disconnect.assert_awaited_once()
client.db.connect.assert_awaited_once()
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
client.db.query_raw.assert_awaited_once_with("SELECT 1")
@@ -140,15 +141,19 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
)
client._db_last_reconnect_attempt_ts = 0.0
client._db_reconnect_cooldown_seconds = 10
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.recreate_prisma_client = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
client._start_engine_watcher = AsyncMock()
# Use a counter-based mock to avoid StopIteration when time.time() is called
# more times than expected (varies by Python version / internal code paths).
fake_clock = iter(range(100, 10000))
with patch(
"litellm.proxy.utils.time.time", side_effect=lambda: float(next(fake_clock))
with (
patch(
"litellm.proxy.utils.time.time",
side_effect=lambda: float(next(fake_clock)),
),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
result = await client.attempt_db_reconnect(
reason="unit_test_cooldown_timestamp_after_attempt",
@@ -162,23 +167,28 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
@pytest.mark.asyncio
async def test_run_reconnect_cycle_watchdog_should_use_direct_db_ops(
async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client(
mock_proxy_logging,
):
"""Direct reconnect goes through recreate_prisma_client (which non-blockingly
kills the old engine) instead of calling disconnect() see issue #26191.
"""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.disconnect = AsyncMock(side_effect=AssertionError("wrapper disconnect used"))
client.connect = AsyncMock(side_effect=AssertionError("wrapper connect used"))
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.disconnect = AsyncMock(
side_effect=AssertionError("disconnect must not be called")
)
client.db.recreate_prisma_client = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
client._start_engine_watcher = AsyncMock()
await client._run_reconnect_cycle(timeout_seconds=None)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await client._run_reconnect_cycle(timeout_seconds=None)
client.db.disconnect.assert_awaited_once()
client.db.connect.assert_awaited_once()
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
client.db.query_raw.assert_awaited_once_with("SELECT 1")
client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
@@ -189,19 +199,22 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client._db_watchdog_reconnect_timeout_seconds = 0.1
client.db.disconnect = AsyncMock(return_value=None)
client._start_engine_watcher = AsyncMock()
async def _slow_connect():
async def _slow_recreate(_db_url):
await asyncio.sleep(0.08)
async def _slow_query(_query: str):
await asyncio.sleep(0.08)
return [{"result": 1}]
client.db.connect = AsyncMock(side_effect=_slow_connect)
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
client.db.query_raw = AsyncMock(side_effect=_slow_query)
with pytest.raises(asyncio.TimeoutError):
with (
pytest.raises(asyncio.TimeoutError),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await client._run_reconnect_cycle(timeout_seconds=None)
@@ -212,19 +225,22 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget(
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.disconnect = AsyncMock(return_value=None)
client._start_engine_watcher = AsyncMock()
async def _slow_connect():
async def _slow_recreate(_db_url):
await asyncio.sleep(0.08)
async def _slow_query(_query: str):
await asyncio.sleep(0.08)
return [{"result": 1}]
client.db.connect = AsyncMock(side_effect=_slow_connect)
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
client.db.query_raw = AsyncMock(side_effect=_slow_query)
with pytest.raises(asyncio.TimeoutError):
with (
pytest.raises(asyncio.TimeoutError),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await client._run_reconnect_cycle(timeout_seconds=0.1)
@@ -319,42 +335,154 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
@pytest.mark.asyncio
async def test_lightweight_reconnect_kills_engine_on_disconnect_failure(
async def test_recreate_prisma_client_kills_old_engine_without_disconnect(
mock_proxy_logging,
):
"""Lightweight reconnect must kill the old engine PID when disconnect() fails."""
"""recreate_prisma_client SIGTERMs the old engine PID directly rather than
calling `disconnect()`, which blocks the asyncio event loop on the sync
`subprocess.Popen.wait()` inside prisma-client-py see issue #26191.
"""
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}])
disconnect_mock = AsyncMock(
side_effect=AssertionError("disconnect must not be called on reconnect path")
)
client.db._original_prisma.disconnect = disconnect_mock
with (
patch.object(client, "_get_engine_pid", return_value=9999),
patch("os.kill") as mock_kill,
patch("asyncio.sleep", new_callable=AsyncMock),
patch.object(client.db, "_get_engine_pid", return_value=9999),
patch("litellm.proxy.db.prisma_client.os.kill") as mock_kill,
patch("litellm.proxy.db.prisma_client.asyncio.sleep", new_callable=AsyncMock),
):
await client._run_reconnect_cycle(timeout_seconds=5.0)
# Return a Prisma instance whose connect() is awaitable.
fake_new_prisma = MagicMock()
fake_new_prisma.connect = AsyncMock(return_value=None)
with patch("prisma.Prisma", return_value=fake_new_prisma):
await client.db.recreate_prisma_client("postgresql://test")
mock_kill.assert_any_call(9999, signal.SIGTERM)
client.db.connect.assert_awaited_once()
client.db.query_raw.assert_awaited_once_with("SELECT 1")
disconnect_mock.assert_not_awaited()
fake_new_prisma.connect.assert_awaited_once()
# ---------------------------------------------------------------------------
# get_generic_data: transport-reconnect-and-retry coverage (issue #25143)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(
async def test_get_generic_data_retries_on_transport_error_for_config_table(
mock_proxy_logging,
):
"""Lightweight reconnect must NOT kill when disconnect() succeeds."""
"""`get_generic_data(table_name="config")` self-heals on a transient
`httpx.ReadError`: reconnect once, retry once, return the row.
Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry
branch that 1.82.6 had on this method. `_update_config_from_db` fans out
four concurrent `get_generic_data` calls, so a single transport flap used
to surface as four `db_exceptions` alerts and a stale config window.
"""
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)
expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}}
invocations: list[None] = []
mock_kill.assert_not_called()
async def _flaky_find_first(**kwargs):
invocations.append(None)
if len(invocations) == 1:
raise httpx.ReadError("simulated transport blip")
return expected_row
client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first)
client.attempt_db_reconnect = AsyncMock(return_value=True)
result = await client.get_generic_data(
key="param_name",
value="general_settings",
table_name="config",
)
assert result == expected_row
assert len(invocations) == 2
client.attempt_db_reconnect.assert_awaited_once()
reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs
assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
# The failure_handler telemetry side-effect must NOT fire on the first
# transport blip — only if the post-retry call also fails. Drain the
# event loop so any spuriously-spawned task would have run by now.
await asyncio.sleep(0)
mock_proxy_logging.failure_handler.assert_not_called()
@pytest.mark.asyncio
async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging):
"""If reconnect itself does not succeed, propagate the original transport
error and let the existing failure_handler / db_exceptions telemetry fire."""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.litellm_config.find_first = AsyncMock(
side_effect=httpx.ReadError("simulated transport blip")
)
client.attempt_db_reconnect = AsyncMock(return_value=False)
with pytest.raises(httpx.ReadError):
await client.get_generic_data(
key="param_name",
value="general_settings",
table_name="config",
)
client.attempt_db_reconnect.assert_awaited_once()
# Failure telemetry IS expected here — the read genuinely failed.
await asyncio.sleep(0)
mock_proxy_logging.failure_handler.assert_called_once()
# ---------------------------------------------------------------------------
# _engine_confirmed_dead flag-reset bug (B2)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
mock_proxy_logging,
):
"""Regression test for the flag-reset bug.
Before the fix, `_run_reconnect_cycle` cleared
`self._engine_confirmed_dead = False` *before* awaiting
`_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout,
missing DATABASE_URL, recreate failure), the flag was left cleared and the
next attempt could demote to the lightweight path even though the engine
was genuinely dead.
The fix moves the reset into the success branch the flag must stay True
when heavy reconnect raises.
"""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client._engine_confirmed_dead = True
client._engine_pid = 0 # so `_is_engine_alive` is not consulted
# Make the heavy reconnect path raise.
client.db.recreate_prisma_client = AsyncMock(
side_effect=RuntimeError("simulated heavy reconnect failure")
)
client._start_engine_watcher = AsyncMock()
client._cleanup_engine_watcher = MagicMock()
client._reap_all_zombies = MagicMock()
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
with pytest.raises(Exception):
await client._run_reconnect_cycle(timeout_seconds=5.0)
# The flag must STILL be True so the next attempt re-enters the heavy
# branch instead of silently demoting to the lightweight path.
assert client._engine_confirmed_dead is True
@@ -0,0 +1,611 @@
"""
Unit tests for workflow management endpoints (/v1/workflows/runs/*).
Uses FastAPI TestClient with a mocked prisma_client.
"""
import os
import sys
from datetime import datetime, timezone
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import FastAPI
from fastapi.testclient import TestClient
from prisma.errors import UniqueViolationError
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.management_endpoints.workflow_management_endpoints import router
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_run(
run_id: str = "run-1",
session_id: str = "sess-1",
workflow_type: str = "shin-builder",
status: str = "pending",
created_by: Any = "tok-test",
) -> MagicMock:
obj = MagicMock()
obj.run_id = run_id
obj.session_id = session_id
obj.workflow_type = workflow_type
obj.status = status
obj.created_by = created_by
obj.created_at = datetime.now(timezone.utc)
obj.updated_at = datetime.now(timezone.utc)
obj.input = None
obj.output = None
obj.metadata = None
return obj
def _make_event(
event_id: str = "evt-1",
run_id: str = "run-1",
event_type: str = "step.started",
step_name: str = "grill",
sequence_number: int = 0,
) -> MagicMock:
obj = MagicMock()
obj.event_id = event_id
obj.run_id = run_id
obj.event_type = event_type
obj.step_name = step_name
obj.sequence_number = sequence_number
obj.data = None
obj.created_at = datetime.now(timezone.utc)
return obj
def _make_message(
message_id: str = "msg-1",
run_id: str = "run-1",
role: str = "user",
content: str = "hello",
sequence_number: int = 0,
) -> MagicMock:
obj = MagicMock()
obj.message_id = message_id
obj.run_id = run_id
obj.role = role
obj.content = content
obj.sequence_number = sequence_number
obj.session_id = None
obj.created_at = datetime.now(timezone.utc)
return obj
def _make_tx(event_return=None, run_return=None, msg_return=None) -> MagicMock:
"""Build an async context-manager mock for prisma_client.db.tx()."""
tx = MagicMock()
tx.litellm_workflowevent = MagicMock()
tx.litellm_workflowevent.create = AsyncMock(
return_value=event_return or _make_event()
)
tx.litellm_workflowrun = MagicMock()
tx.litellm_workflowrun.update = AsyncMock(return_value=run_return or _make_run())
tx.litellm_workflowmessage = MagicMock()
tx.litellm_workflowmessage.create = AsyncMock(
return_value=msg_return or _make_message()
)
tx.__aenter__ = AsyncMock(return_value=tx)
tx.__aexit__ = AsyncMock(return_value=False)
return tx
def _make_prisma_client() -> MagicMock:
client = MagicMock()
client.db = MagicMock()
client.db.litellm_workflowrun = MagicMock()
client.db.litellm_workflowevent = MagicMock()
client.db.litellm_workflowmessage = MagicMock()
# default tx() returns a no-op transaction
client.db.tx = MagicMock(return_value=_make_tx())
return client
def _make_app() -> FastAPI:
app = FastAPI()
app.include_router(router)
return app
def _override_auth() -> Any:
from litellm.proxy._types import UserAPIKeyAuth
auth = UserAPIKeyAuth(api_key="sk-test", user_id="admin")
auth.token = "tok-test"
return auth
def _override_auth_admin() -> Any:
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
auth = UserAPIKeyAuth(api_key="sk-master")
auth.user_role = LitellmUserRoles.PROXY_ADMIN # type: ignore[assignment]
return auth
def _override_auth_user_with_token(token: str = "tok-abc") -> Any:
"""Return a non-admin caller whose hashed token equals `token`."""
from litellm.proxy._types import UserAPIKeyAuth
auth = UserAPIKeyAuth(api_key="sk-user", user_id="user-1")
auth.token = token # override the computed hash with a predictable value
return auth
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestCreateWorkflowRun:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_create_returns_run(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.create = AsyncMock(return_value=_make_run())
resp = self.client.post(
"/v1/workflows/runs",
json={"workflow_type": "shin-builder"},
)
assert resp.status_code == 200
self._prisma.db.litellm_workflowrun.create.assert_awaited_once()
@patch("litellm.proxy.proxy_server.prisma_client", None)
def test_create_500_when_no_db(self):
resp = self.client.post(
"/v1/workflows/runs",
json={"workflow_type": "shin-builder"},
)
assert resp.status_code == 500
class TestListWorkflowRuns:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_returns_runs(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(
return_value=[_make_run()]
)
resp = self.client.get("/v1/workflows/runs")
assert resp.status_code == 200
data = resp.json()
assert data["count"] == 1
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_filters_by_status(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
resp = self.client.get("/v1/workflows/runs?status=running")
assert resp.status_code == 200
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
assert call_kwargs["where"]["status"] == "running"
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_filters_by_multiple_statuses(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
resp = self.client.get("/v1/workflows/runs?status=running,paused")
assert resp.status_code == 200
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
assert call_kwargs["where"]["status"] == {"in": ["running", "paused"]}
class TestGetWorkflowRun:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_get_existing_run(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
resp = self.client.get("/v1/workflows/runs/run-1")
assert resp.status_code == 200
@patch("litellm.proxy.proxy_server.prisma_client")
def test_get_missing_run_returns_404(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
resp = self.client.get("/v1/workflows/runs/nonexistent")
assert resp.status_code == 404
class TestUpdateWorkflowRun:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_update_status(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
updated = _make_run(status="completed")
self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=updated)
resp = self.client.patch(
"/v1/workflows/runs/run-1", json={"status": "completed"}
)
assert resp.status_code == 200
self._prisma.db.litellm_workflowrun.update.assert_awaited_once()
@patch("litellm.proxy.proxy_server.prisma_client")
def test_update_no_fields_returns_400(self, mock_pc):
mock_pc.db = self._prisma.db
resp = self.client.patch("/v1/workflows/runs/run-1", json={})
assert resp.status_code == 400
class TestAppendWorkflowEvent:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_append_event_updates_run_status(self, mock_pc):
mock_pc.db = self._prisma.db
# _require_run check
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
tx = _make_tx(
event_return=_make_event(), run_return=_make_run(status="running")
)
self._prisma.db.tx = MagicMock(return_value=tx)
resp = self.client.post(
"/v1/workflows/runs/run-1/events",
json={"event_type": "step.started", "step_name": "grill"},
)
assert resp.status_code == 200
# run status updated inside tx
tx.litellm_workflowrun.update.assert_awaited_once()
update_call = tx.litellm_workflowrun.update.call_args[1]
assert update_call["data"]["status"] == "running"
@patch("litellm.proxy.proxy_server.prisma_client")
def test_append_event_no_status_update_for_unknown_type(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
tx = _make_tx(event_return=_make_event(event_type="custom.event"))
self._prisma.db.tx = MagicMock(return_value=tx)
resp = self.client.post(
"/v1/workflows/runs/run-1/events",
json={"event_type": "custom.event", "step_name": "grill"},
)
assert resp.status_code == 200
# no status update inside tx for unknown event_type
tx.litellm_workflowrun.update.assert_not_awaited()
@patch("litellm.proxy.proxy_server.prisma_client")
def test_sequence_number_increments(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
existing = _make_event(sequence_number=4)
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(
return_value=[existing]
)
tx = _make_tx(event_return=_make_event(sequence_number=5))
self._prisma.db.tx = MagicMock(return_value=tx)
self.client.post(
"/v1/workflows/runs/run-1/events",
json={"event_type": "step.started", "step_name": "plan"},
)
create_call = tx.litellm_workflowevent.create.call_args[1]
assert create_call["data"]["sequence_number"] == 5
@patch("litellm.proxy.proxy_server.prisma_client")
def test_unknown_run_id_returns_404(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
resp = self.client.post(
"/v1/workflows/runs/nonexistent/events",
json={"event_type": "step.started", "step_name": "grill"},
)
assert resp.status_code == 404
@patch("litellm.proxy.proxy_server.prisma_client")
def test_sequence_collision_retries_and_succeeds(self, mock_pc):
"""UniqueViolationError on first attempt triggers retry; second attempt succeeds."""
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
# First tx raises UniqueViolationError; second succeeds.
tx_fail = _make_tx()
tx_fail.__aenter__ = AsyncMock(return_value=tx_fail)
tx_fail.litellm_workflowevent.create = AsyncMock(
side_effect=UniqueViolationError(
{"user_facing_error": {"message": "unique"}}
)
)
tx_fail.__aexit__ = AsyncMock(return_value=False)
tx_ok = _make_tx(event_return=_make_event(sequence_number=1))
self._prisma.db.tx = MagicMock(side_effect=[tx_fail, tx_ok])
resp = self.client.post(
"/v1/workflows/runs/run-1/events",
json={"event_type": "step.started", "step_name": "grill"},
)
assert resp.status_code == 200
class TestWorkflowMessages:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_append_message(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
self._prisma.db.litellm_workflowmessage.create = AsyncMock(
return_value=_make_message()
)
resp = self.client.post(
"/v1/workflows/runs/run-1/messages",
json={"role": "user", "content": "fix the bug"},
)
assert resp.status_code == 200
@patch("litellm.proxy.proxy_server.prisma_client")
def test_append_message_unknown_run_returns_404(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
resp = self.client.post(
"/v1/workflows/runs/nonexistent/messages",
json={"role": "user", "content": "hello"},
)
assert resp.status_code == 404
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_messages_ordered(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(
return_value=[
_make_message(sequence_number=0),
_make_message(sequence_number=1, role="assistant"),
]
)
resp = self.client.get("/v1/workflows/runs/run-1/messages")
assert resp.status_code == 200
data = resp.json()
assert data["count"] == 2
call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1]
assert call_kwargs["order"] == {"sequence_number": "asc"}
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_messages_respects_limit(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
resp = self.client.get("/v1/workflows/runs/run-1/messages?limit=25")
assert resp.status_code == 200
call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1]
assert call_kwargs["take"] == 25
class TestListWorkflowEvents:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_events_ordered(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(
return_value=[
_make_event(sequence_number=0),
_make_event(sequence_number=1),
]
)
resp = self.client.get("/v1/workflows/runs/run-1/events")
assert resp.status_code == 200
data = resp.json()
assert data["count"] == 2
call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1]
assert call_kwargs["order"] == {"sequence_number": "asc"}
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_events_respects_limit(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run()
)
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
resp = self.client.get("/v1/workflows/runs/run-1/events?limit=10")
assert resp.status_code == 200
call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1]
assert call_kwargs["take"] == 10
@patch("litellm.proxy.proxy_server.prisma_client")
def test_list_events_unknown_run_returns_404(self, mock_pc):
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
resp = self.client.get("/v1/workflows/runs/nonexistent/events")
assert resp.status_code == 404
class TestTenantIsolation:
"""Ownership enforcement: non-admin callers only see their own runs."""
def _make_app_with_auth(self, auth_fn):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
self._prisma = _make_prisma_client()
app = _make_app()
app.dependency_overrides[user_api_key_auth] = auth_fn
return TestClient(app, raise_server_exceptions=True)
@patch("litellm.proxy.proxy_server.prisma_client")
def test_create_stores_caller_token(self, mock_pc):
token = "tok-owner"
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.create = AsyncMock(
return_value=_make_run(created_by=token)
)
resp = client.post("/v1/workflows/runs", json={"workflow_type": "test"})
assert resp.status_code == 200
create_call = self._prisma.db.litellm_workflowrun.create.call_args[1]
assert create_call["data"]["created_by"] == token
@patch("litellm.proxy.proxy_server.prisma_client")
def test_non_admin_list_scoped_to_caller_token(self, mock_pc):
token = "tok-owner"
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
resp = client.get("/v1/workflows/runs")
assert resp.status_code == 200
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
assert call_kwargs["where"].get("created_by") == token
@patch("litellm.proxy.proxy_server.prisma_client")
def test_admin_list_not_scoped(self, mock_pc):
client = self._make_app_with_auth(_override_auth_admin)
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
resp = client.get("/v1/workflows/runs")
assert resp.status_code == 200
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
assert "created_by" not in call_kwargs["where"]
@patch("litellm.proxy.proxy_server.prisma_client")
def test_non_admin_get_other_users_run_returns_404(self, mock_pc):
token = "tok-caller"
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
mock_pc.db = self._prisma.db
# Run owned by a different key
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run(created_by="tok-other-owner")
)
resp = client.get("/v1/workflows/runs/run-1")
assert resp.status_code == 404
@patch("litellm.proxy.proxy_server.prisma_client")
def test_non_admin_get_null_owner_run_returns_404(self, mock_pc):
token = "tok-caller"
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run(created_by=None)
)
resp = client.get("/v1/workflows/runs/run-1")
assert resp.status_code == 404
@patch("litellm.proxy.proxy_server.prisma_client")
def test_non_admin_update_null_owner_run_returns_404(self, mock_pc):
token = "tok-caller"
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run(created_by=None)
)
self._prisma.db.litellm_workflowrun.update = AsyncMock(
return_value=_make_run(status="completed")
)
resp = client.patch("/v1/workflows/runs/run-1", json={"status": "completed"})
assert resp.status_code == 404
self._prisma.db.litellm_workflowrun.update.assert_not_awaited()
@patch("litellm.proxy.proxy_server.prisma_client")
def test_non_admin_get_own_run_succeeds(self, mock_pc):
token = "tok-caller"
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
mock_pc.db = self._prisma.db
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
return_value=_make_run(created_by=token)
)
resp = client.get("/v1/workflows/runs/run-1")
assert resp.status_code == 200
@@ -218,6 +218,141 @@ class TestProxyBaseLLMRequestProcessing:
headers_with_invalid
)
@pytest.mark.asyncio
async def test_build_litellm_proxy_success_headers_from_llm_response(self):
"""
Google native :generateContent uses this helper instead of base_process_llm_request;
ensure x-litellm-* headers and callback hooks merge like the main proxy path.
"""
mock_request = MagicMock(spec=Request)
mock_request.headers = {}
class _FakeGenaiResponse:
_hidden_params = {
"model_id": "deployment-model-id",
"cache_key": "ck-test",
"api_base": "https://generativelanguage.googleapis.com/v1beta",
"response_cost": 0.001,
"additional_headers": {"llm_provider-ratelimit-requests": "1000"},
}
logging_obj = MagicMock()
logging_obj.litellm_call_id = "call-id-test"
mock_user = MagicMock()
mock_user.tpm_limit = None
mock_user.rpm_limit = None
mock_user.max_budget = None
mock_user.spend = 0.0
mock_user.allowed_model_region = None
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
return_value={"x-ratelimit-remaining-requests": "999"}
)
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=_FakeGenaiResponse(),
request_data={"model": "gemini/gemini-1.5-flash"},
request=mock_request,
user_api_key_dict=mock_user,
logging_obj=logging_obj,
version="9.9.9",
proxy_logging_obj=proxy_logging_obj,
)
assert headers["x-litellm-call-id"] == "call-id-test"
assert headers["x-litellm-model-id"] == "deployment-model-id"
assert headers["x-litellm-version"] == "9.9.9"
assert headers["llm_provider-ratelimit-requests"] == "1000"
assert headers["x-ratelimit-remaining-requests"] == "999"
proxy_logging_obj.post_call_response_headers_hook.assert_awaited_once()
@pytest.mark.asyncio
async def test_build_litellm_proxy_success_headers_streaming_style_iterator(self):
"""AsyncGoogleGenAIGenerateContentStreamingIterator sets _hidden_params at init; headers must propagate."""
class _FakeStreamLike:
def __aiter__(self):
return self
async def __anext__(self):
raise StopAsyncIteration
_hidden_params = {
"model_id": "stream-model-id",
"api_base": "https://generativelanguage.googleapis.com/v1beta",
"cache_key": "",
"response_cost": "",
"additional_headers": {"llm_provider-x": "y"},
}
mock_request = MagicMock(spec=Request)
mock_request.headers = {}
logging_obj = MagicMock()
logging_obj.litellm_call_id = "cid-stream"
mock_user = MagicMock()
mock_user.tpm_limit = None
mock_user.rpm_limit = None
mock_user.max_budget = None
mock_user.spend = 0.0
mock_user.allowed_model_region = None
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=_FakeStreamLike(),
request_data={"model": "gemini/gemini-2.0-flash"},
request=mock_request,
user_api_key_dict=mock_user,
logging_obj=logging_obj,
version="1.0.0",
proxy_logging_obj=proxy_logging_obj,
)
assert headers["x-litellm-model-id"] == "stream-model-id"
assert headers["x-litellm-model-api-base"] == (
"https://generativelanguage.googleapis.com/v1beta"
)
assert headers["llm_provider-x"] == "y"
@pytest.mark.asyncio
async def test_build_litellm_proxy_success_headers_no_hidden_params_metadata_fallback(
self,
):
"""When response has no _hidden_params, model_id can still come from litellm_metadata."""
class _BareResponse:
pass
mock_request = MagicMock(spec=Request)
mock_request.headers = {}
logging_obj = MagicMock()
logging_obj.litellm_call_id = "cid-meta"
mock_user = MagicMock()
mock_user.tpm_limit = None
mock_user.rpm_limit = None
mock_user.max_budget = None
mock_user.spend = 0.0
mock_user.allowed_model_region = None
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=_BareResponse(),
request_data={
"model": "gemini/gemini-1.5-flash",
"litellm_metadata": {"model_info": {"id": "meta-model-id"}},
},
request=mock_request,
user_api_key_dict=mock_user,
logging_obj=logging_obj,
version="1.0.0",
proxy_logging_obj=proxy_logging_obj,
)
assert headers["x-litellm-model-id"] == "meta-model-id"
@pytest.mark.asyncio
async def test_add_litellm_data_to_request_with_stream_timeout_header(self):
"""
@@ -1158,13 +1293,6 @@ class TestCommonRequestProcessingHelpers:
assert mock_tracer.trace.call_count == 4
# Verify that each call was made with the correct operation name
expected_calls = [
(("streaming.chunk.yield",), {}),
(("streaming.chunk.yield",), {}),
(("streaming.chunk.yield",), {}),
(("streaming.chunk.yield",), {}),
]
actual_calls = mock_tracer.trace.call_args_list
assert len(actual_calls) == 4
@@ -2544,6 +2544,14 @@ class TestPriceDataReloadAPI:
class TestPriceDataReloadIntegration:
"""Integration tests for the complete price data reload feature"""
@pytest.fixture(autouse=True)
def _flush_litellm_config_cache(self):
from litellm.proxy.utils import litellm_config_cache
litellm_config_cache.flush_cache()
yield
litellm_config_cache.flush_cache()
@pytest.fixture
def client_with_auth(self):
"""Create a test client with authentication"""
@@ -2601,6 +2609,7 @@ class TestPriceDataReloadIntegration:
def test_distributed_reload_check_function(self):
"""Test the _check_and_reload_model_cost_map function"""
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import litellm_config_cache
proxy_config = ProxyConfig()
@@ -2609,14 +2618,19 @@ class TestPriceDataReloadIntegration:
# Test case 1: No config in database
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
# _check_and_reload_model_cost_map routes through get_config_param,
# which calls prisma.get_generic_data on a cache miss.
mock_prisma.get_generic_data = AsyncMock(return_value=None)
# Should return early without reloading
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Test case 2: Config with interval but not time to reload
litellm_config_cache.flush_cache()
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": 6, "force_reload": False}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
# Mock current time and last reload time
with patch(
@@ -2632,8 +2646,10 @@ class TestPriceDataReloadIntegration:
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Test case 3: Config with force reload
litellm_config_cache.flush_cache()
mock_config.param_value = {"interval_hours": 6, "force_reload": True}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
original_model_cost = litellm.model_cost.copy()
@@ -2675,6 +2691,8 @@ class TestPriceDataReloadIntegration:
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": 24, "force_reload": True}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
# _check_and_reload_model_cost_map now reads through get_generic_data.
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
original_model_cost = litellm.model_cost.copy()
@@ -2770,6 +2788,8 @@ class TestPriceDataReloadIntegration:
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": 12, "force_reload": True}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
# _check_and_reload_anthropic_beta_headers now reads through get_generic_data.
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
with patch(
@@ -29,8 +29,21 @@ from litellm.types.utils import (
)
@pytest.fixture(autouse=True)
def _use_local_model_cost_map(monkeypatch):
original_model_cost = litellm.model_cost
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.get_model_info.cache_clear()
try:
yield
finally:
litellm.model_cost = original_model_cost
litellm.get_model_info.cache_clear()
class TestGPTImageCostCalculator:
"""Test the OpenAI gpt-image-1 cost calculator"""
"""Test the OpenAI gpt-image cost calculator"""
def test_gpt_image_1_cost_with_text_only(self):
"""Test cost calculation with only text input tokens"""
@@ -149,6 +162,44 @@ class TestGPTImageCostCalculator:
assert cost == 0.0
def test_gpt_image_2_cost_with_text_and_image_tokens(self):
"""Test cost calculation for gpt-image-2 token pricing"""
from litellm.llms.openai.image_generation.cost_calculator import cost_calculator
usage = Usage(
prompt_tokens=600,
completion_tokens=5000,
total_tokens=5600,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=100,
image_tokens=500,
),
completion_tokens_details=CompletionTokensDetailsWrapper(
text_tokens=1000,
image_tokens=4000,
),
)
image_response = ImageResponse(
created=1234567890,
data=[ImageObject(url="http://example.com/image.jpg")],
)
image_response.usage = usage
cost = cost_calculator(
model="gpt-image-2",
image_response=image_response,
custom_llm_provider="openai",
)
# GPT Image 2 pricing:
# Text input: 100 * $5/1M = 0.0005
# Image input: 500 * $8/1M = 0.004
# Text output: 1000 * $10/1M = 0.01
# Image output: 4000 * $30/1M = 0.12
expected_cost = 0.0005 + 0.004 + 0.01 + 0.12
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
class TestGPTImageCostRouting:
"""Test that gpt-image models are properly routed to the token-based calculator"""
@@ -182,6 +233,33 @@ class TestGPTImageCostRouting:
expected_cost = 0.0005 + 0.2
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
def test_openai_gpt_image_2_routes_to_token_calculator(self):
"""Test that OpenAI gpt-image-2 routes to token-based calculator"""
from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
usage = Usage(
prompt_tokens=100,
completion_tokens=5000,
total_tokens=5100,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100),
completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=5000),
)
image_response = ImageResponse(
created=1234567890,
data=[ImageObject(url="http://example.com/image.jpg")],
)
image_response.usage = usage
cost = CostCalculatorUtils.route_image_generation_cost_calculator(
model="gpt-image-2",
completion_response=image_response,
custom_llm_provider="openai",
)
expected_cost = 0.0005 + 0.15
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
def test_openai_dalle_routes_to_pixel_calculator(self):
"""Test that OpenAI DALL-E still routes to pixel-based calculator"""
from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
+93 -14
View File
@@ -32,6 +32,19 @@ from litellm.utils import (
# Adds the parent directory to the system path
@pytest.fixture
def local_model_cost_map(monkeypatch):
original_model_cost = litellm.model_cost
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.get_model_info.cache_clear()
try:
yield
finally:
litellm.model_cost = original_model_cost
litellm.get_model_info.cache_clear()
def test_check_provider_match_azure_ai_allows_openai_and_azure():
"""
Test that azure_ai provider can match openai and azure models.
@@ -198,6 +211,72 @@ def test_get_optional_params_image_gen_filters_empty_values():
assert optional_params == {}
def test_gpt_image_provider_detection_covers_existing_family():
for image_model in ("gpt-image-1", "gpt-image-1-mini", "gpt-image-1.5"):
model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=image_model)
assert model == image_model
assert custom_llm_provider == "openai"
def test_gpt_image_2_provider_and_model_info(local_model_cost_map):
model, custom_llm_provider, _, _ = litellm.get_llm_provider(model="gpt-image-2")
assert model == "gpt-image-2"
assert custom_llm_provider == "openai"
model_info = litellm.get_model_info(model="gpt-image-2")
assert model_info["litellm_provider"] == "openai"
assert model_info["mode"] == "image_generation"
assert model_info["input_cost_per_token"] == 5e-06
assert model_info["input_cost_per_image_token"] == 8e-06
assert model_info["output_cost_per_token"] == 1e-05
assert model_info["output_cost_per_image_token"] == 3e-05
assert (
"/v1/images/generations"
in litellm.model_cost["gpt-image-2"]["supported_endpoints"]
)
assert (
"/v1/images/edits" in litellm.model_cost["gpt-image-2"]["supported_endpoints"]
)
assert model_info["supports_vision"] is True
assert model_info["supports_pdf_input"] is True
def test_gpt_image_2_snapshot_model_info(local_model_cost_map):
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model="gpt-image-2-2026-04-21"
)
assert model == "gpt-image-2-2026-04-21"
assert custom_llm_provider == "openai"
model_info = litellm.get_model_info(model="gpt-image-2-2026-04-21")
assert model_info["litellm_provider"] == "openai"
assert model_info["mode"] == "image_generation"
assert model_info["output_cost_per_image_token"] == 3e-05
def test_azure_gpt_image_2_model_info(local_model_cost_map):
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model="azure/gpt-image-2"
)
assert model == "gpt-image-2"
assert custom_llm_provider == "azure"
model_info = litellm.get_model_info(
model="gpt-image-2", custom_llm_provider="azure"
)
assert model_info["litellm_provider"] == "azure"
assert model_info["mode"] == "image_generation"
assert model_info["input_cost_per_token"] == 5e-06
assert model_info["input_cost_per_image_token"] == 8e-06
assert model_info["output_cost_per_token"] == 1e-05
assert model_info["output_cost_per_image_token"] == 3e-05
def test_all_model_configs():
from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import (
VertexAIAi21Config,
@@ -1179,7 +1258,7 @@ def test_get_model_info_shows_supports_computer_use():
"model, custom_llm_provider",
[
("gpt-3.5-turbo", "openai"),
("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"),
("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"),
("gemini-2.5-pro", "vertex_ai"),
],
)
@@ -1325,7 +1404,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/bedrock-claude-3-opus",
"bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
),
(
@@ -1623,7 +1702,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/bedrock-claude-3-opus",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"Bedrock Claude 3 Opus via Converse API",
),
@@ -1710,7 +1789,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/staging-claude-opus",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"Staging Claude Opus",
),
@@ -1722,7 +1801,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/high-performance-claude",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"High-performance Claude deployment",
),
@@ -1860,7 +1939,7 @@ class TestProxyFunctionCalling:
bedrock_models = [
"bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0",
"bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
]
for model in bedrock_models:
@@ -1892,7 +1971,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/bedrock-claude-3-opus",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"Bedrock Claude 3 Opus via Converse API",
),
@@ -1979,7 +2058,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/staging-claude-opus",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"Staging Claude Opus",
),
@@ -1991,7 +2070,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/high-performance-claude",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"High-performance Claude deployment",
),
@@ -2129,7 +2208,7 @@ class TestProxyFunctionCalling:
bedrock_models = [
"bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0",
"bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
]
for model in bedrock_models:
@@ -2161,7 +2240,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/bedrock-claude-3-opus",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"Bedrock Claude 3 Opus via Converse API",
),
@@ -2248,7 +2327,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/staging-claude-opus",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"Staging Claude Opus",
),
@@ -2260,7 +2339,7 @@ class TestProxyFunctionCalling:
),
(
"litellm_proxy/high-performance-claude",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
False,
"High-performance Claude deployment",
),
@@ -2398,7 +2477,7 @@ class TestProxyFunctionCalling:
bedrock_models = [
"bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0",
"bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0",
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
]
for model in bedrock_models:
@@ -267,6 +267,7 @@ class TestRAGFlowVectorStore(BaseVectorStoreTest):
api_base="http://localhost:9380",
litellm_logging_obj=logging_obj,
litellm_params={},
extra_body=None,
)
def test_transform_search_vector_store_response_not_implemented(self):