diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 34ae3638a5b..6115a444cee 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -392,6 +392,7 @@ class DualCache(BaseCache): value: float, parent_otel_span: Optional[Span] = None, local_only: bool = False, + refresh_ttl: bool = False, **kwargs, ) -> Optional[float]: """ @@ -399,6 +400,9 @@ class DualCache(BaseCache): Value - float - the value you want to increment by + Refresh_ttl - bool - if True, resets the Redis TTL on every write. + Default False preserves window-style semantics. + Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ @@ -415,6 +419,7 @@ class DualCache(BaseCache): value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), + refresh_ttl=refresh_ttl, ) return result diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 84a2887f527..deee4f6ea48 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -824,6 +824,7 @@ class RedisCache(BaseCache): value: float, ttl: Optional[int] = None, parent_otel_span: Optional[Span] = None, + refresh_ttl: bool = False, ) -> float: from redis.asyncio import Redis @@ -834,11 +835,12 @@ class RedisCache(BaseCache): try: result = await _redis_client.incrbyfloat(name=key, amount=value) if _used_ttl is not None: - # check if key already has ttl, if not -> set ttl - current_ttl = await _redis_client.ttl(key) - if current_ttl == -1: - # Key has no expiration + if refresh_ttl: await _redis_client.expire(key, _used_ttl) + else: + current_ttl = await _redis_client.ttl(key) + if current_ttl == -1: + await _redis_client.expire(key, _used_ttl) ## LOGGING ## end_time = time.time() diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a34b73b5313..bc2ca805e94 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2535,10 +2535,16 @@ class BaseLLMHTTPHandler: }, ) + delete_kwargs: Dict[str, Any] = { + "url": url, + "headers": headers, + "timeout": timeout, + } + if data: + delete_kwargs["json"] = data + try: - response = await async_httpx_client.delete( - url=url, headers=headers, json=data, timeout=timeout - ) + response = await async_httpx_client.delete(**delete_kwargs) except Exception as e: raise self._handle_error( @@ -2619,10 +2625,16 @@ class BaseLLMHTTPHandler: }, ) + delete_kwargs: Dict[str, Any] = { + "url": url, + "headers": headers, + "timeout": timeout, + } + if data: + delete_kwargs["json"] = data + try: - response = sync_httpx_client.delete( - url=url, headers=headers, json=data, timeout=timeout - ) + response = sync_httpx_client.delete(**delete_kwargs) except Exception as e: raise self._handle_error( diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index e486336cec0..0928ce914da 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -52,6 +52,37 @@ class ResetBudgetJob: ### RESET MULTI-WINDOW BUDGETS ### await self.reset_budget_windows() + @staticmethod + async def _invalidate_spend_counter(counter_key: str) -> None: + """Zero a spend counter so a DB-row reset takes effect immediately. + + Call AFTER the DB write commits. Clearing Redis before the DB + commit opens a window where get_current_spend reads 0 from Redis + while the DB still holds the pre-reset value, allowing bypass. + """ + try: + from litellm.proxy.proxy_server import spend_counter_cache + + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=0.0, ttl=60 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0, ttl=60 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset spend counter %s in Redis: %s. " + "Budget may be over-enforced until counter expires.", + counter_key, + redis_err, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to reset spend counter %s: %s", counter_key, e + ) + async def reset_budget_for_litellm_team_members( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -64,46 +95,30 @@ class ResetBudgetJob: if budget.budget_id is not None ] - # Reset spend counters for affected team members. - # Reset Redis directly so a transient failure doesn't leave stale - # counters that get_current_spend would read as authoritative. try: - from litellm.proxy.proxy_server import spend_counter_cache - memberships = await self.prisma_client.db.litellm_teammembership.find_many( where={"budget_id": {"in": budget_ids}} ) - for m in memberships: - counter_key = f"spend:team_member:{m.user_id}:{m.team_id}" - # Always reset in-memory - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 - ) - # Explicitly reset Redis with warning on failure - if spend_counter_cache.redis_cache is not None: - try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 - ) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to reset team member spend counter in Redis %s: %s. " - "Budget may be over-enforced until counter expires.", - counter_key, - redis_err, - ) except Exception as e: + memberships = [] verbose_proxy_logger.warning( - "Failed to reset team member spend counters: %s", e + "Failed to fetch team memberships for counter invalidation: %s", e ) - return await self.prisma_client.db.litellm_teammembership.update_many( + update_result = await self.prisma_client.db.litellm_teammembership.update_many( where={"budget_id": {"in": budget_ids}}, data={ "spend": 0, }, ) + for m in memberships: + await self._invalidate_spend_counter( + f"spend:team_member:{m.user_id}:{m.team_id}" + ) + + return update_result + async def reset_budget_for_keys_linked_to_budgets( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -126,17 +141,36 @@ class ResetBudgetJob: if not budget_ids: return - return await self.prisma_client.db.litellm_verificationtoken.update_many( - where={ - "budget_id": {"in": budget_ids}, - "budget_duration": None, # only keys without their own reset schedule - "spend": {"gt": 0}, # only reset keys that have accumulated spend - }, - data={ - "spend": 0, - }, + where_clause: dict = { + "budget_id": {"in": budget_ids}, + "budget_duration": None, # only keys without their own reset schedule + "spend": {"gt": 0}, # only reset keys that have accumulated spend + } + + try: + keys = await self.prisma_client.db.litellm_verificationtoken.find_many( + where=where_clause + ) + except Exception as e: + keys = [] + verbose_proxy_logger.warning( + "Failed to fetch keys for counter invalidation: %s", e + ) + + update_result = ( + await self.prisma_client.db.litellm_verificationtoken.update_many( + where=where_clause, + data={ + "spend": 0, + }, + ) ) + for k in keys: + await self._invalidate_spend_counter(f"spend:key:{k.token}") + + return update_result + async def reset_budget_for_litellm_budget_table(self): """ Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired @@ -365,6 +399,10 @@ class ResetBudgetJob: data_list=updated_keys, table_name="key", ) + for k in updated_keys: + token = getattr(k, "token", None) + if token: + await self._invalidate_spend_counter(f"spend:key:{token}") end_time = time.time() if len(failed_keys) > 0: # If any keys failed to reset @@ -450,6 +488,12 @@ class ResetBudgetJob: data_list=updated_users, table_name="user", ) + for u in updated_users: + user_id = getattr(u, "user_id", None) + if user_id: + await self._invalidate_spend_counter( + f"spend:user:{user_id}" + ) end_time = time.time() if len(failed_users) > 0: # If any users failed to reset @@ -541,6 +585,12 @@ class ResetBudgetJob: data_list=updated_teams, table_name="team", ) + for t in updated_teams: + team_id = getattr(t, "team_id", None) + if team_id: + await self._invalidate_spend_counter( + f"spend:team:{team_id}" + ) end_time = time.time() if len(failed_teams) > 0: # If any teams failed to reset diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index bf60a087c65..a979471dc8e 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -129,7 +129,9 @@ class SpendCounterReseed: """ lock = await SpendCounterReseed._get_lock(counter_key) async with lock: - # Re-check after acquiring the lock - another waiter may have warmed it. + # Re-check after acquiring the lock. Skip in-memory on a clean + # Redis miss - in-memory is per-pod-stale. + redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: val = await spend_counter_cache.redis_cache.async_get_cache( @@ -137,11 +139,13 @@ class SpendCounterReseed: ) if val is not None: return float(val) + redis_clean_miss = True except Exception: pass - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + if not redis_clean_miss: + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key) if db_spend is None: @@ -149,7 +153,7 @@ class SpendCounterReseed: # Warm even when 0 so subsequent reads hit cache, not DB. try: await spend_counter_cache.async_increment_cache( - key=counter_key, value=db_spend + key=counter_key, value=db_spend, refresh_ttl=True ) except Exception: verbose_proxy_logger.exception( diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 6ada8f58783..967ac9f0ac4 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,10 +1,6 @@ -from datetime import datetime +from fastapi import APIRouter, Depends, Request, Response +from fastapi.responses import ORJSONResponse -from fastapi import APIRouter, Depends, HTTPException, Request, Response -from fastapi.responses import ORJSONResponse, StreamingResponse - -import litellm -from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -30,12 +26,17 @@ async def google_generate_content( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ( general_settings, llm_router, proxy_config, proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, version, ) @@ -43,48 +44,33 @@ async def google_generate_content( if "model" not in data: data["model"] = model_name - # Extract generationConfig and pass it as config parameter - generation_config = data.pop("generationConfig", None) - if generation_config: - data["config"] = generation_config - - # Add user authentication metadata for cost tracking - data = await add_litellm_data_to_request( - data=data, - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - general_settings=general_settings, - version=version, - ) - - # Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id - data["litellm_call_id"] = request.headers.get( - "x-litellm-call-id", str(uuid.uuid4()) - ) - logging_obj, data = litellm.utils.function_setup( - original_function="agenerate_content", - rules_obj=litellm.utils.Rules(), - start_time=datetime.now(), - **data, - ) - data["litellm_logging_obj"] = logging_obj - - # call router - if llm_router is None: - raise HTTPException(status_code=500, detail="Router not initialized") - response = await llm_router.agenerate_content(**data) - success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( - response=response, - request_data=data, - request=request, - user_api_key_dict=user_api_key_dict, - logging_obj=logging_obj, - version=version, - proxy_logging_obj=proxy_logging_obj, - ) - fastapi_response.headers.update(success_headers) - return response + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="agenerate_content", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=model_name, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) @router.post( @@ -101,73 +87,52 @@ async def google_stream_generate_content( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ( general_settings, llm_router, proxy_config, proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, version, ) data = await _read_request_body(request=request) - if "model" not in data: data["model"] = model_name + data["stream"] = True - data["stream"] = True # enforce streaming for this endpoint - - # Extract generationConfig and pass it as config parameter - generation_config = data.pop("generationConfig", None) - if generation_config: - data["config"] = generation_config - - # Add user authentication metadata for cost tracking - data = await add_litellm_data_to_request( - data=data, - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - general_settings=general_settings, - version=version, - ) - - # Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id - data["litellm_call_id"] = request.headers.get( - "x-litellm-call-id", str(uuid.uuid4()) - ) - logging_obj, data = litellm.utils.function_setup( - original_function="agenerate_content_stream", - rules_obj=litellm.utils.Rules(), - start_time=datetime.now(), - **data, - ) - data["litellm_logging_obj"] = logging_obj - - # call router - if llm_router is None: - raise HTTPException(status_code=500, detail="Router not initialized") - response = await llm_router.agenerate_content_stream(**data) - - success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( - response=response, - request_data=data, - request=request, - user_api_key_dict=user_api_key_dict, - logging_obj=logging_obj, - version=version, - proxy_logging_obj=proxy_logging_obj, - ) - - # Check if response is an async iterator (streaming response) - if response is not None and hasattr(response, "__aiter__"): - return StreamingResponse( - content=response, - media_type="text/event-stream", - headers=success_headers, + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="agenerate_content_stream", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=model_name, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, ) - fastapi_response.headers.update(success_headers) - return response @router.post( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6f5ab1afb68..6cba6a3e96b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1798,12 +1798,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: 3. Reseed from authoritative DB spend (counter expired, cross-pod stale) 4. Caller-supplied fallback (DB unavailable, cold start) """ - # 1. Try Redis first (cross-pod authoritative) + # 1. Redis first (cross-pod authoritative). On clean miss, skip + # in-memory: per-pod in-memory only has this pod's writes, so it + # would mask cross-pod increments. + redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: return float(val) + redis_clean_miss = True except Exception as e: verbose_proxy_logger.debug( "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", @@ -1811,10 +1815,11 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: e, ) - # 2. Fall back to in-memory counter (single-instance or Redis failure) - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + # 2. In-memory only when Redis is unreachable. + if not redis_clean_miss: + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass. db_spend = await SpendCounterReseed.coalesced( @@ -1976,10 +1981,12 @@ async def _init_and_increment_spend_counter( base_spend = getattr(source, "spend", 0.0) or 0.0 if base_spend > 0: await spend_counter_cache.async_increment_cache( - key=counter_key, value=base_spend + key=counter_key, value=base_spend, refresh_ttl=True ) - await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) + await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, refresh_ttl=True + ) async def update_cache( # noqa: PLR0915 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index b39eb42821c..78192400fb0 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -50,6 +50,50 @@ async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping ) +@pytest.mark.asyncio +async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl( + monkeypatch, redis_no_ping +): + """With refresh_ttl=True, every increment should call expire() to bump + the TTL, even when the key already has a TTL (counter-style use).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment( + key="spend:team_member:u:t", value=0.05, refresh_ttl=True + ) + + mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60) + + +@pytest.mark.asyncio +async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( + monkeypatch, redis_no_ping +): + """Default (refresh_ttl=False) preserves window-style semantics: TTL is + set only on first creation, never refreshed (used by rate-limit windows).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment(key="rate_limit:window", value=1) + + mock_redis_instance.expire.assert_not_awaited() + + @pytest.mark.asyncio async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "my-fake-host") diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 752b5ff0905..b846cd600f0 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,3 +1,4 @@ +import asyncio import os import sys from unittest.mock import AsyncMock, Mock, patch @@ -8,6 +9,8 @@ import pytest sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, _google_genai_streaming_hidden_params, @@ -103,7 +106,9 @@ def test_fingerprint_agentic_tools_is_deterministic(): tools_a = {"tool_calls": [{"id": "1", "input": {"q": "abc"}, "name": "web_search"}]} tools_b = {"tool_calls": [{"name": "web_search", "input": {"q": "abc"}, "id": "1"}]} - assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b) + assert handler._fingerprint_agentic_tools( + tools_a + ) == handler._fingerprint_agentic_tools(tools_b) @pytest.mark.asyncio @@ -350,3 +355,70 @@ def test_google_genai_streaming_hidden_params_model_info_and_router_fallback(): response_headers=httpx.Headers({}), ) assert from_router["model_id"] == "router-model-id" + + +def _build_delete_response_mock(captured: dict): + """Returns a fake httpx delete that records its kwargs.""" + + def _response() -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id": "resp_x", "object": "response", "deleted": true}', + request=httpx.Request(method="DELETE", url="https://test.openai.azure.com"), + ) + + async def fake_async_delete(*args, **kwargs): + captured.update(kwargs) + return _response() + + def fake_sync_delete(*args, **kwargs): + captured.update(kwargs) + return _response() + + return fake_async_delete, fake_sync_delete + + +def test_async_delete_responses_omits_body_for_azure(): + """Azure responses DELETE rejects requests with any body. Verify the handler + does not pass `json=` to httpx when the transformer returns an empty dict.""" + captured: dict = {} + fake_async_delete, _ = _build_delete_response_mock(captured) + + async def run(): + with patch.object(AsyncHTTPHandler, "delete", new=fake_async_delete): + await litellm.adelete_responses( + response_id="resp_xyz", + custom_llm_provider="azure", + api_base="https://test.openai.azure.com", + api_key="test-key", + api_version="2025-03-01-preview", + ) + + asyncio.run(run()) + + assert "json" not in captured + assert "data" not in captured + assert captured["url"].endswith( + "/openai/responses/resp_xyz?api-version=2025-03-01-preview" + ) + + +def test_sync_delete_responses_omits_body_for_azure(): + captured: dict = {} + _, fake_sync_delete = _build_delete_response_mock(captured) + + with patch.object(HTTPHandler, "delete", new=fake_sync_delete): + litellm.delete_responses( + response_id="resp_xyz", + custom_llm_provider="azure", + api_base="https://test.openai.azure.com", + api_key="test-key", + api_version="2025-03-01-preview", + ) + + assert "json" not in captured + assert "data" not in captured + assert captured["url"].endswith( + "/openai/responses/resp_xyz?api-version=2025-03-01-preview" + ) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 379ccf4d9af..5c86f9057a1 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1049,3 +1049,159 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch): asyncio.run(job.reset_budget_windows()) # must not raise prisma_client.db.litellm_teamtable.update.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Counter invalidation on budget reset +# --------------------------------------------------------------------------- + + +def _make_counter_invalidation_job(monkeypatch): + """Stub spend_counter_cache so we can observe invalidation calls.""" + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = MagicMock() + spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + return spend_counter_cache + + +def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch): + """Team-member budget reset clears the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + expired_budget = type("B", (), {"budget_id": "budget-1"}) + membership = type( + "Membership", + (), + {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"}, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_teammembership.find_many = AsyncMock( + return_value=[membership] + ) + prisma_client.db.litellm_teammembership.update_many = AsyncMock( + return_value={"count": 1} + ) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:team_member:alice:team-x", value=0.0, ttl=60 + ) + counter_cache.redis_cache.async_set_cache.assert_any_await( + key="spend:team_member:alice:team-x", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_keys_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """Key budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["key"] = [ + type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": now, + "id": "key-1", + "token": "sk-abc", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:key:sk-abc", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_users_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """User budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["user"] = [ + type( + "User", + (), + { + "spend": 50.0, + "budget_duration": "7d", + "budget_reset_at": now, + "id": "user-1", + "user_id": "alice", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:user:alice", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_teams_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """Team budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["team"] = [ + type( + "Team", + (), + { + "spend": 200.0, + "budget_duration": "1mo", + "budget_reset_at": now, + "id": "team-1", + "team_id": "team-x", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:team:team-x", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch): + """Resetting keys via budget tier must clear each linked key's counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + expired_budget = type("B", (), {"budget_id": "budget-1"}) + linked_key = type("Key", (), {"token": "sk-linked"}) + + prisma_client = MagicMock() + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[linked_key] + ) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 1} + ) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:key:sk-linked", value=0.0, ttl=60 + ) diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index a35f358f365..434f7953c21 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -4,7 +4,7 @@ Test to verify the Google GenAI proxy API endpoints """ import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -13,520 +13,171 @@ sys.path.insert( ) # Adds the parent directory to the system path -def test_google_generate_content_endpoint(): - """Test that the google_generate_content endpoint correctly routes requests""" - # Skip this test if we can't import the required modules due to missing dependencies - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient +def _build_test_client(): + from fastapi import FastAPI + from fastapi.testclient import TestClient - from litellm.proxy.google_endpoints.endpoints import router as google_router + from litellm.proxy.google_endpoints.endpoints import router as google_router + + app = FastAPI() + app.include_router(google_router) + return TestClient(app) + + +def _patch_base_process(return_value=None): + """Patch ProxyBaseLLMRequestProcessing.base_process_llm_request so endpoint + tests don't run the full pipeline. Returns the AsyncMock so callers can + inspect call args.""" + if return_value is None: + return_value = {"test": "response"} + return patch( + "litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new_callable=AsyncMock, + return_value=return_value, + ) + + +def test_google_generate_content_endpoint(): + """generateContent routes through ProxyBaseLLMRequestProcessing with the + agenerate_content route_type — that pipeline runs pre_call_hook + + during_call_hook + post_call_success_hook for every guardrail callback.""" + try: + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock the router's agenerate_content method - with patch("litellm.proxy.proxy_server.llm_router") as mock_router: - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Send a request to the endpoint + with _patch_base_process() as mock_base: response = client.post( "/v1beta/models/test-model:generateContent", json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - # Verify the response assert response.status_code == 200 - assert response.json() == {"test": "response"} - - # Verify that agenerate_content was called - mock_router.agenerate_content.assert_called_once() + mock_base.assert_called_once() + kwargs = mock_base.call_args.kwargs + assert kwargs["route_type"] == "agenerate_content" + assert kwargs["model"] == "test-model" def test_google_stream_generate_content_endpoint(): - """Test that the google_stream_generate_content endpoint correctly routes streaming requests""" - # Skip this test if we can't import the required modules due to missing dependencies + """streamGenerateContent must route through the same processor with the + streaming route_type so the guardrail pipeline runs.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock the router's agenerate_content_stream method to return a stream - async def mock_stream_generator(): - yield 'data: {"test": "stream_chunk_1"}\n\n' - yield 'data: {"test": "stream_chunk_2"}\n\n' - yield "data: [DONE]\n\n" - - with patch("litellm.proxy.proxy_server.llm_router") as mock_router: - mock_router.agenerate_content_stream = AsyncMock( - return_value=mock_stream_generator() - ) - - # Send a request to the endpoint + with ( + _patch_base_process() as mock_base, + patch( + "litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__", + return_value=None, + ) as mock_init, + ): response = client.post( "/v1beta/models/test-model:streamGenerateContent", json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - # Verify the response assert response.status_code == 200 + mock_base.assert_called_once() + kwargs = mock_base.call_args.kwargs + assert kwargs["route_type"] == "agenerate_content_stream" + assert kwargs["model"] == "test-model" - # Verify that agenerate_content_stream was called with correct parameters - mock_router.agenerate_content_stream.assert_called_once() - call_args = mock_router.agenerate_content_stream.call_args - assert call_args[1]["stream"] is True - assert call_args[1]["model"] == "test-model" - assert call_args[1]["contents"] == [ + # stream=True must be forced into the data the processor receives. + init_kwargs = mock_init.call_args.kwargs + assert init_kwargs["data"]["stream"] is True + assert init_kwargs["data"]["model"] == "test-model" + assert init_kwargs["data"]["contents"] == [ {"role": "user", "parts": [{"text": "Hello"}]} ] -def test_google_generate_content_with_cost_tracking_metadata(): - """Test that the google_generate_content endpoint includes user metadata for cost tracking""" +def test_google_generate_content_data_flows_through_processor(): + """The body the client sends must reach ProxyBaseLLMRequestProcessing + intact so the pipeline can apply guardrails to it.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), + _patch_base_process(), patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, + "litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__", + return_value=None, + ) as mock_init, ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - # Simulate adding user metadata - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - "user_api_key_team_id": "test-team-id", - "user_api_key": "hashed-key", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request to the endpoint - response = client.post( + client.post( "/v1beta/models/test-model:generateContent", - json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, - headers={"Authorization": "Bearer sk-test-key"}, - ) - - # Verify the response - assert response.status_code == 200 - - # Verify that add_litellm_data_to_request was called - mock_add_data.assert_called_once() - - # Verify that agenerate_content was called with metadata - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that litellm_metadata exists and contains user information - assert "litellm_metadata" in called_data - assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" - assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id" - - -def test_google_stream_generate_content_with_cost_tracking_metadata(): - """Test that the google_stream_generate_content endpoint includes user metadata for cost tracking""" - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router - except ImportError as e: - pytest.skip(f"Skipping test due to missing dependency: {e}") - - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock the router's agenerate_content_stream method to return a stream - mock_stream = AsyncMock() - mock_stream.__aiter__ = lambda self: mock_stream - mock_stream.__anext__.side_effect = StopAsyncIteration - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream) - - # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - # Simulate adding user metadata - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - "user_api_key_team_id": "test-team-id", - "user_api_key": "hashed-key", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request to the endpoint - response = client.post( - "/v1beta/models/test-model:streamGenerateContent", - json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, - headers={"Authorization": "Bearer sk-test-key"}, - ) - - # Verify the response - assert response.status_code == 200 - - # Verify that add_litellm_data_to_request was called - mock_add_data.assert_called_once() - - # Verify that agenerate_content_stream was called with metadata - mock_router.agenerate_content_stream.assert_called_once() - call_args = mock_router.agenerate_content_stream.call_args - called_data = call_args[1] - - # Verify that litellm_metadata exists and contains user information - assert "litellm_metadata" in called_data - assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" - assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id" - # Verify stream is set to True - assert called_data["stream"] is True - - -def test_google_generate_content_with_system_instruction(): - """ - Test that systemInstruction is correctly passed through from the endpoint to the router. - - This test verifies the fix for systemInstruction being dropped when forwarding - requests to Vertex AI through the Google GenAI endpoint. - """ - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router - except ImportError as e: - pytest.skip(f"Skipping test due to missing dependency: {e}") - - # Create a FastAPI app and include the router - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Define the systemInstruction to test - system_instruction = {"parts": [{"text": "Your name is Doodle."}]} - - # Send a request with systemInstruction - response = client.post( - "/v1beta/models/gemini-2.5-pro:generateContent", json={ - "systemInstruction": system_instruction, - "contents": [ - {"parts": [{"text": "What is your name?"}], "role": "user"} - ], - }, - headers={"Authorization": "Bearer sk-test-key"}, - ) - - # Verify the response - assert response.status_code == 200 - - # Verify that agenerate_content was called - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that systemInstruction is present in the call arguments - assert "systemInstruction" in called_data - assert called_data["systemInstruction"] == system_instruction - assert ( - called_data["systemInstruction"]["parts"][0]["text"] - == "Your name is Doodle." - ) - - # Verify contents are also present - assert "contents" in called_data - assert len(called_data["contents"]) == 1 - assert called_data["contents"][0]["role"] == "user" - - -def test_google_generate_content_with_image_config(): - """ - Test that imageConfig is correctly passed through from generationConfig to the router. - - This test verifies that imageConfig parameters (aspectRatio, imageSize) are preserved - when forwarding requests to Google GenAI through the endpoint. - """ - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router - except ImportError as e: - pytest.skip(f"Skipping test due to missing dependency: {e}") - - # Create a FastAPI app and include the router - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request with generationConfig containing imageConfig - response = client.post( - "/v1beta/models/gemini-3-pro-image-preview:generateContent", - json={ - "contents": [ - { - "role": "user", - "parts": [ - { - "text": "Create a vibrant infographic about photosynthesis" - } - ], - } - ], + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "systemInstruction": {"parts": [{"text": "Your name is Doodle."}]}, "generationConfig": { "responseModalities": ["TEXT", "IMAGE"], "imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"}, }, }, - headers={"Authorization": "Bearer sk-test-key"}, ) - # Verify the response - assert response.status_code == 200 - - # Verify that agenerate_content was called - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that config is present in the call arguments - assert "config" in called_data - - # Verify that imageConfig is preserved in the config - assert "imageConfig" in called_data["config"] - assert called_data["config"]["imageConfig"]["aspectRatio"] == "9:16" - assert called_data["config"]["imageConfig"]["imageSize"] == "4K" - - # Verify that responseModalities is also preserved - assert "responseModalities" in called_data["config"] - assert called_data["config"]["responseModalities"] == ["TEXT", "IMAGE"] - - # Verify contents are also present - assert "contents" in called_data - assert len(called_data["contents"]) == 1 - assert called_data["contents"][0]["role"] == "user" + data = mock_init.call_args.kwargs["data"] + assert data["model"] == "test-model" + assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + assert data["systemInstruction"] == { + "parts": [{"text": "Your name is Doodle."}] + } + # generationConfig arrives intact here; the rename to `config` is + # done downstream in route_request (see test_route_llm_request). + assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + assert data["generationConfig"]["imageConfig"]["aspectRatio"] == "9:16" -def test_google_generate_content_metadata_and_trace_id_callbacks(): - """Test that google_generate_content sets litellm_call_id and logging_obj for callbacks (e.g. S3, Langfuse)""" +def test_google_generate_content_forwards_call_id_header(): + """The endpoint must forward the x-litellm-call-id header to the processor + so the helper can stamp it on the logging object. Trace continuity from + client → callbacks (S3, Langfuse, etc.) depends on this header surviving + the hop through these endpoints.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - # Simulate adding user metadata - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request to the endpoint with x-litellm-call-id header - test_call_id = "test-custom-call-id" - response = client.post( + with _patch_base_process() as mock_base: + client.post( "/v1beta/models/test-model:generateContent", json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, - headers={ - "Authorization": "Bearer sk-test-key", - "x-litellm-call-id": test_call_id, - }, + headers={"x-litellm-call-id": "trace-abc-123"}, ) - assert response.status_code == 200 - - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that the litellm_logging_obj got assigned in the final called_data to router - assert "litellm_logging_obj" in called_data - assert "litellm_call_id" in called_data - assert called_data["litellm_call_id"] == test_call_id + forwarded_request = mock_base.call_args.kwargs["request"] + assert forwarded_request.headers.get("x-litellm-call-id") == "trace-abc-123" -def test_google_stream_generate_content_metadata_and_trace_id_callbacks(): - """Test that google_stream_generate_content sets litellm_call_id and logging_obj for callbacks""" +def test_google_count_tokens_unchanged(): + """countTokens has its own path and isn't affected by the pipeline change.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - app = FastAPI() - app.include_router(google_router) - client = TestClient(app) + fake_response = MagicMock() + fake_response.original_response = { + "totalTokens": 7, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], + } + fake_response.total_tokens = 7 - mock_stream = AsyncMock() - mock_stream.__aiter__ = lambda self: mock_stream - mock_stream.__anext__.side_effect = StopAsyncIteration - - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + return_value=fake_response, ): - mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream) - - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - test_call_id = "test-custom-stream-call-id" response = client.post( - "/v1beta/models/test-model:streamGenerateContent", - json={"contents": [{"role": "user", "parts": [{"text": "Hello stream"}]}]}, - headers={ - "Authorization": "Bearer sk-test-key", - "x-litellm-call-id": test_call_id, - }, + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) assert response.status_code == 200 - - mock_router.agenerate_content_stream.assert_called_once() - call_args = mock_router.agenerate_content_stream.call_args - called_data = call_args[1] - - assert "litellm_logging_obj" in called_data - assert "litellm_call_id" in called_data - assert called_data["litellm_call_id"] == test_call_id + body = response.json() + assert body["totalTokens"] == 7 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e0b6d229e2c..37e53005650 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5750,3 +5750,90 @@ class TestLazyFeatureMiddleware: assert attempts == [ "called" ], f"failing register_fn should be invoked once, not on every request; got {attempts}" + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory(): + """When Redis is reachable and cleanly returns None (TTL expired, + counter genuinely absent), the read must reseed from DB - NOT fall + through to per-pod in-memory which only contains this pod's writes. + + Pre-fix in multi-pod deployments, in-memory contained a stale local + subset (e.g. $30) while DB had the true cross-pod total ($500). The + fall-through returned $30, enforcement passed, bypass. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + + # Per-pod stale in-memory: only this pod's writes, not cross-pod truth. + counter_cache.in_memory_cache.set_cache(key=counter_key, value=30.0) + + # Redis cleanly returns None (key expired or never written on this pod). + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_increment = AsyncMock(return_value=500.0) + counter_cache.redis_cache = fake_redis + + # DB has the authoritative cross-pod spend. + db_row = MagicMock() + db_row.spend = 500.0 + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + assert spend == 500.0, ( + f"expected DB-authoritative 500.0 on clean Redis miss, got {spend} " + f"(stale per-pod in-memory $30 would have caused multi-pod bypass)" + ) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_error_falls_back_to_in_memory(): + """When Redis raises, the read should still degrade to in-memory rather + than going straight to DB - in-memory is at least same-pod-fresh and + cheaper than a DB query during a Redis outage.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + + counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0) + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( + return_value=MagicMock(spend=999.0) + ) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + assert spend == 42.0, ( + f"expected in-memory fallback 42.0 on Redis error, got {spend} " + f"(should not have hit DB when Redis errored)" + ) + # DB query should NOT have fired - in-memory short-circuits. + fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 96870b6cc77..bfea21e705e 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -239,3 +239,55 @@ async def test_route_request_with_router_settings_override_preserves_existing(): assert call_kwargs["num_retries"] == 10 # Key/team timeout should be applied since not in request assert call_kwargs["timeout"] == 30 + + +@pytest.mark.parametrize( + "route_type", ["agenerate_content", "agenerate_content_stream"] +) +@pytest.mark.asyncio +async def test_route_request_maps_generation_config_for_google_routes(route_type): + """For Google generate_content routes, route_request must rename + `generationConfig` (Google's wire format) to `config` (the kwarg the + router method expects). Without this mapping the request reaches the + LLM with the field under the wrong name and the config is dropped.""" + data = { + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "generationConfig": { + "responseModalities": ["TEXT", "IMAGE"], + "imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"}, + }, + } + llm_router = MagicMock() + getattr(llm_router, route_type).return_value = "ok" + + await route_request(data, llm_router, None, route_type) + + call_kwargs = getattr(llm_router, route_type).call_args[1] + assert "generationConfig" not in call_kwargs + assert "config" in call_kwargs + assert call_kwargs["config"]["responseModalities"] == ["TEXT", "IMAGE"] + assert call_kwargs["config"]["imageConfig"]["aspectRatio"] == "9:16" + assert call_kwargs["config"]["imageConfig"]["imageSize"] == "4K" + + +@pytest.mark.parametrize( + "route_type", ["agenerate_content", "agenerate_content_stream"] +) +@pytest.mark.asyncio +async def test_route_request_preserves_existing_config_for_google_routes(route_type): + """If the caller already supplies `config`, route_request must not + overwrite it with `generationConfig`.""" + data = { + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "config": {"existing": True}, + "generationConfig": {"shouldNotWin": True}, + } + llm_router = MagicMock() + getattr(llm_router, route_type).return_value = "ok" + + await route_request(data, llm_router, None, route_type) + + call_kwargs = getattr(llm_router, route_type).call_args[1] + assert call_kwargs["config"] == {"existing": True}