diff --git a/docker/Dockerfile.dev b/docker/Dockerfile.dev index e0ecff31ca..fb84230adc 100644 --- a/docker/Dockerfile.dev +++ b/docker/Dockerfile.dev @@ -15,6 +15,7 @@ USER root # Install build dependencies in one layer RUN apt-get update && apt-get install -y --no-install-recommends \ gcc \ + g++ \ python3-dev \ libssl-dev \ pkg-config \ diff --git a/litellm/constants.py b/litellm/constants.py index a9facabb01..1af53b2dae 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1344,6 +1344,7 @@ LITELLM_UI_SESSION_DURATION = os.getenv("LITELLM_UI_SESSION_DURATION", "24h") ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" +DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics" CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data" CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( @@ -1397,6 +1398,10 @@ APSCHEDULER_REPLACE_EXISTING = os.getenv( "1", ] # always replace existing jobs +# The number of tag entries are higher than number of user, team entries. This leads to a higher QPS. +# This will run tag spcific tasks at a later time to smooth QPS +DAILY_TAG_SPEND_BATCH_MULTIPLIER = 2.3 + DEFAULT_HEALTH_CHECK_INTERVAL = int( os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300) ) # 5 minutes diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 166be712d5..7fc28b68e9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -172,6 +172,34 @@ if MCP_AVAILABLE: mcp_info: Optional[MCPInfo] = None model_config = ConfigDict(arbitrary_types_allowed=True) + def _normalize_resource_contents(contents: list) -> List[ReadResourceContents]: + """Normalize ResourceContents to ReadResourceContents, preserving meta (MCP 1.26.0+).""" + normalized: List[ReadResourceContents] = [] + for content in contents: + meta = getattr(content, "meta", None) + if meta is None and hasattr(content, "model_dump"): + d = content.model_dump() + meta = d.get("meta") + if meta is None: + meta = d.get("_meta") + if isinstance(content, TextResourceContents): + normalized.append( + ReadResourceContents( + content=content.text, + mime_type=content.mimeType, + meta=meta, + ) + ) + elif isinstance(content, BlobResourceContents): + normalized.append( + ReadResourceContents( + content=content.blob, + mime_type=content.mimeType, + meta=meta, + ) + ) + return normalized + ######################################################## ############ Initialize the MCP Server ################# ######################################################## @@ -632,26 +660,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) - normalized_contents: List[ReadResourceContents] = [] - for content in read_resource_result.contents: - if isinstance(content, TextResourceContents): - text_content: TextResourceContents = content - normalized_contents.append( - ReadResourceContents( - content=text_content.text, - mime_type=text_content.mimeType, - ) - ) - elif isinstance(content, BlobResourceContents): - blob_content: BlobResourceContents = content - normalized_contents.append( - ReadResourceContents( - content=blob_content.blob, - mime_type=None, - ) - ) - - return normalized_contents + return _normalize_resource_contents(read_resource_result.contents) ######################################################## ############ End of MCP Server Routes ################## diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a305d5be1e..241b66bc0a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -28,7 +28,7 @@ from typing import ( import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache, RedisCache -from litellm.constants import DB_SPEND_UPDATE_JOB_NAME +from litellm.constants import DB_SPEND_UPDATE_JOB_NAME,DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, @@ -797,7 +797,6 @@ class DBSpendUpdateWriter: daily_org_spend_update_queue=self.daily_org_spend_update_queue, daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue, daily_agent_spend_update_queue=self.daily_agent_spend_update_queue, - daily_tag_spend_update_queue=self.daily_tag_spend_update_queue, ) # Only commit from redis to db if this pod is the leader @@ -814,7 +813,6 @@ class DBSpendUpdateWriter: daily_org_spend_update_transactions, daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions, - daily_tag_spend_update_transactions, ) = ( await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() ) @@ -890,13 +888,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_org_spend_update_transactions, ) - if daily_tag_spend_update_transactions is not None: - await DBSpendUpdateWriter.update_daily_tag_spend( - n_retry_times=n_retry_times, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_tag_spend_update_transactions, - ) if daily_end_user_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_end_user_spend( n_retry_times=n_retry_times, @@ -991,19 +982,7 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_org_spend_update_transactions, ) - ################## Daily Tag Spend Update Transactions ################## - # Aggregate all in memory daily tag spend transactions and commit to db - daily_tag_spend_update_transactions = cast( - Dict[str, DailyTagSpendTransaction], - await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), - ) - - await DBSpendUpdateWriter.update_daily_tag_spend( - n_retry_times=n_retry_times, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_tag_spend_update_transactions, - ) + # NOTE: Daily tag spend is committed by a separate scheduler job. ################## Daily End-User Spend Update Transactions ################## # Aggregate all in memory daily end-user spend transactions and commit to db @@ -1032,10 +1011,75 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_agent_spend_update_transactions, ) - + ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) + async def _commit_daily_tag_spend_to_db( + self, + prisma_client: PrismaClient, + n_retry_times: int, + proxy_logging_obj: ProxyLogging, + ): + """ + Commit only tag spend updates to database. + This is called by a separate scheduler job at a longer interval. + """ + daily_tag_spend_update_transactions = cast( + Dict[str, DailyTagSpendTransaction], + await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), + ) + + if daily_tag_spend_update_transactions: + await DBSpendUpdateWriter.update_daily_tag_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_tag_spend_update_transactions, + ) + + async def _commit_daily_tag_spend_to_db_with_redis( + self, + prisma_client: PrismaClient, + n_retry_times: int, + proxy_logging_obj: ProxyLogging, + ): + """ + Commit daily tag spend updates using Redis buffering. + + This lets the dedicated daily tag scheduler drain both in-memory and + Redis-backed tag transactions. + """ + await self.redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis( + daily_tag_spend_update_queue=self.daily_tag_spend_update_queue, + ) + + if await self.pod_lock_manager.acquire_lock( + cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, + ): + verbose_proxy_logger.debug("acquired lock for daily tag spend updates") + try: + daily_tag_spend_update_transactions = await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + + if daily_tag_spend_update_transactions: + await DBSpendUpdateWriter.update_daily_tag_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_tag_spend_update_transactions, + ) + except Exception as e: + verbose_proxy_logger.error( + "Spend tracking - failed to commit daily tag spend updates from Redis to DB. " + "Data already popped from Redis may be lost. Error: %s\n%s", + str(e), + traceback.format_exc(), + ) + finally: + await self.pod_lock_manager.release_lock( + cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, + ) + async def _flush_tool_discovery_queue( self, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index c51c06df2f..bdca867081 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -131,7 +131,6 @@ class RedisUpdateBuffer: daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, - daily_tag_spend_update_queue: DailySpendUpdateQueue, ): """ Stores the in-memory spend updates to Redis @@ -202,9 +201,6 @@ class RedisUpdateBuffer: daily_agent_spend_update_transactions = ( await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) - daily_tag_spend_update_transactions = ( - await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) verbose_proxy_logger.debug( "ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions @@ -245,11 +241,6 @@ class RedisUpdateBuffer: REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, ), - ( - daily_tag_spend_update_transactions, - REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, - ), ] rpush_list: List[RedisPipelineRpushOperation] = [] @@ -376,22 +367,20 @@ class RedisUpdateBuffer: Optional[Dict[str, DailyOrganizationSpendTransaction]], Optional[Dict[str, DailyEndUserSpendTransaction]], Optional[Dict[str, DailyAgentSpendTransaction]], - Optional[Dict[str, DailyTagSpendTransaction]], ]: """ - Drains all 7 Redis buffer queues in a single pipeline round-trip. + Drains the main 6 Redis buffer queues in a single pipeline round-trip. - Returns a 7-tuple of parsed results in this order: + Returns a 6-tuple of parsed results in this order: 0: DBSpendUpdateTransactions 1: daily user spend 2: daily team spend 3: daily org spend 4: daily end-user spend 5: daily agent spend - 6: daily tag spend """ if self.redis_cache is None: - return None, None, None, None, None, None, None + return None, None, None, None, None, None lpop_list: List[RedisPipelineLpopOperation] = [ RedisPipelineLpopOperation( @@ -417,16 +406,12 @@ class RedisUpdateBuffer: key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ), - RedisPipelineLpopOperation( - key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, - ), ] raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) # Pad with None if pipeline returned fewer results than expected - while len(raw_results) < 7: + while len(raw_results) < 6: raw_results.append(None) # Slot 0: DBSpendUpdateTransactions @@ -436,9 +421,9 @@ class RedisUpdateBuffer: if len(parsed) > 0: db_spend = self._combine_list_of_transactions(parsed) - # Slots 1-6: daily spend categories + # Slots 1-5: daily spend categories daily_results: List[Optional[Dict[str, Any]]] = [] - for slot in range(1, 7): + for slot in range(1, 6): if raw_results[slot] is None: daily_results.append(None) else: @@ -457,7 +442,22 @@ class RedisUpdateBuffer: ), cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]), cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]), - cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]), + ) + + async def store_in_memory_daily_tag_spend_updates_in_redis( + self, + daily_tag_spend_update_queue: DailySpendUpdateQueue, + ) -> None: + """ + Flush in-memory daily tag spend updates and append them to Redis. + """ + daily_tag_spend_update_transactions = ( + await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + await self._store_transactions_in_redis( + transactions=daily_tag_spend_update_transactions, + redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, ) async def get_all_daily_spend_update_transactions_from_redis_buffer( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1b98d1ac84..9738ae4f1a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -54,6 +54,7 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, + DAILY_TAG_SPEND_BATCH_MULTIPLIER ) from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, @@ -6284,6 +6285,25 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + ### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ### + ## Reduces QPS as there are more tags for a single request + tag_spend_update_interval = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER) + from litellm.proxy.utils import update_daily_tag_spend + + scheduler.add_job( + update_daily_tag_spend, + "interval", + seconds=tag_spend_update_interval, + args=[prisma_client, proxy_logging_obj], + id="update_daily_tag_spend_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info( + f"Tag spend update job scheduled at {tag_spend_update_interval}s interval " + f"({tag_spend_update_interval / batch_writing_interval:.1f}x main job interval)" + ) + ### MONITOR SPEND LOGS QUEUE (queue-size-based job) ### if general_settings.get("disable_spend_logs", False) is False: from litellm.proxy.utils import _monitor_spend_logs_queue diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 34c22d2dee..ec98cfd4d1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4827,6 +4827,9 @@ async def update_spend( # noqa: PLR0915 Triggered every minute. + NOTE: This job now skips tag spend updates, which are handled by a separate + scheduler job (update_daily_tag_spend) at a longer interval to reduce contention. + Requires: user_id_list: dict, keys_list: list, @@ -4859,6 +4862,46 @@ async def update_spend( # noqa: PLR0915 ) +async def update_daily_tag_spend( + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, +): + """ + Separate scheduler job to commit daily tag spend updates. + + Runs at a longer interval (2.3x default) than the main update_spend job + to reduce query contention for DailyTagSpend table. + + This is called by a dedicated scheduler job and does NOT process: + - Regular spend updates (user, key, team, org) + - End-user spend + - Agent spend + - Spend logs + + Only processes tag spend transactions from the daily_tag_spend_update_queue. + + Args: + prisma_client: PrismaClient instance + proxy_logging_obj: ProxyLogging instance for error handling + """ + n_retry_times = 3 + try: + if proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis(): + await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + ) + else: + await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.error(f"Error updating daily tag spend: {e}") + + async def update_spend_logs_job( prisma_client: PrismaClient, db_writer_client: Optional[AsyncHTTPHandler], diff --git a/tests/proxy_unit_tests/test_update_daily_tag_spend.py b/tests/proxy_unit_tests/test_update_daily_tag_spend.py new file mode 100644 index 0000000000..7ceeedadae --- /dev/null +++ b/tests/proxy_unit_tests/test_update_daily_tag_spend.py @@ -0,0 +1,134 @@ +from typing import Dict +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.utils import update_daily_tag_spend +from litellm.proxy._types import DailyTagSpendTransaction +import httpx +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + +@pytest.mark.asyncio +async def test_update_daily_tag_spend_delegates_to_tag_commit_writer(): + prisma_client = MagicMock() + proxy_logging_obj = MagicMock() + redis_update_buffer = MagicMock() + redis_update_buffer._should_commit_spend_updates_to_redis.return_value = False + proxy_logging_obj.db_spend_update_writer = MagicMock() + proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock() + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock() + + await update_daily_tag_spend( + prisma_client, + proxy_logging_obj, + ) + + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_awaited_once_with( + prisma_client=prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging_obj, + ) + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis.assert_not_awaited() + +@pytest.mark.asyncio +async def test_update_daily_tag_spend_logs_error_and_does_not_raise(): + prisma_client = MagicMock() + proxy_logging_obj = MagicMock() + redis_update_buffer = MagicMock() + redis_update_buffer._should_commit_spend_updates_to_redis.return_value = False + proxy_logging_obj.db_spend_update_writer = MagicMock() + proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock( + side_effect=ValueError("boom") + ) + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock() + + with patch("litellm.proxy.utils.verbose_proxy_logger.error") as error_logger: + await update_daily_tag_spend( + prisma_client, + proxy_logging_obj, + ) + + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_awaited_once() + error_logger.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_daily_tag_spend_uses_redis_writer_when_enabled(): + prisma_client = MagicMock() + proxy_logging_obj = MagicMock() + redis_update_buffer = MagicMock() + redis_update_buffer._should_commit_spend_updates_to_redis.return_value = True + proxy_logging_obj.db_spend_update_writer = MagicMock() + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock() + proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock() + + await update_daily_tag_spend( + prisma_client, + proxy_logging_obj, + ) + + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis.assert_awaited_once_with( + prisma_client=prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging_obj, + ) + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_daily_tag_spend_retries_then_succeeds(): + prisma_client = MagicMock() + proxy_logging_obj = MagicMock() + + mock_batcher = MagicMock() + mock_table = MagicMock() + mock_batcher.litellm_dailytagspend = mock_table + + # Fail entering batch context 3 times with retryable DB errors, then succeed. + prisma_client.db.batch_.return_value.__aenter__ = AsyncMock( + side_effect=[ + httpx.ConnectError("x"), + httpx.ConnectError("x"), + httpx.ConnectError("x"), + mock_batcher, + ] + ) + + daily_spend_transactions: Dict[str, DailyTagSpendTransaction] = { + "k": { + "tag": "prod-tag", + "date": "2026-04-03", + "api_key": "key-1", + "model": "gpt-4o", + "model_group": None, + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": "", + "endpoint": "", + "prompt_tokens": 10, + "completion_tokens": 5, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "spend": 0.01, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + "request_id": None, + } + } + + with patch("asyncio.sleep", new_callable=AsyncMock) as sleep_mock, patch( + "random.uniform", return_value=0 + ): + await DBSpendUpdateWriter.update_daily_tag_spend( + n_retry_times=3, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_spend_transactions, + ) + + assert prisma_client.db.batch_.return_value.__aenter__.await_count == 4 + assert sleep_mock.await_count == 3 + mock_table.upsert.assert_called_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 8cab568956..384d428888 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,12 +1,11 @@ import asyncio from datetime import datetime, timedelta -from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from mcp import ReadResourceResult, Resource -from mcp.types import Prompt, ResourceTemplate, TextResourceContents +from mcp.types import BlobResourceContents, Prompt, ResourceTemplate, TextResourceContents from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -413,6 +412,111 @@ async def test_mcp_read_resource_success(): assert result is read_result +def test_normalize_resource_contents_passes_metadata(): + """Test that _normalize_resource_contents preserves meta from ResourceContents (MCP 1.26.0+).""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _normalize_resource_contents, + ) + except ImportError: + pytest.skip("MCP server not available") + + meta = {"version": "1.0", "source": "test"} + contents = [ + TextResourceContents( + uri="https://example.com/resource", + text="hello world", + mimeType="text/plain", + meta=meta, + ) + ] + + result = _normalize_resource_contents(contents) + + assert len(result) == 1 + assert result[0].content == "hello world" + assert result[0].mime_type == "text/plain" + assert result[0].meta == meta + + +def test_normalize_resource_contents_blob_with_metadata(): + """Test that _normalize_resource_contents preserves meta for BlobResourceContents.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _normalize_resource_contents, + ) + except ImportError: + pytest.skip("MCP server not available") + + meta = {"encoding": "base64"} + contents = [ + BlobResourceContents( + uri="https://example.com/image.png", + blob="aGVsbG8=", + mimeType="image/png", + meta=meta, + ) + ] + + result = _normalize_resource_contents(contents) + + assert len(result) == 1 + assert result[0].content == "aGVsbG8=" + assert result[0].mime_type == "image/png" + assert result[0].meta == meta + + +def test_normalize_resource_contents_preserves_empty_metadata(): + """Test that empty dict meta is preserved (truthiness bug fix).""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _normalize_resource_contents, + ) + except ImportError: + pytest.skip("MCP server not available") + + empty_meta: dict = {} + contents = [ + TextResourceContents( + uri="https://example.com/resource", + text="hi", + mimeType="text/plain", + meta=empty_meta, + ) + ] + + result = _normalize_resource_contents(contents) + + assert len(result) == 1 + assert result[0].meta == empty_meta + assert result[0].meta is not None + assert result[0].meta == {} + + +def test_normalize_resource_contents_without_metadata(): + """Test that _normalize_resource_contents works when meta is absent (backward compat).""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _normalize_resource_contents, + ) + except ImportError: + pytest.skip("MCP server not available") + + contents = [ + TextResourceContents( + uri="https://example.com/resource", + text="hello", + mimeType="text/plain", + ) + ] + + result = _normalize_resource_contents(contents) + + assert len(result) == 1 + assert result[0].content == "hello" + assert result[0].meta is None + + @pytest.mark.asyncio async def test_mcp_read_resource_multiple_servers_error(): try: @@ -707,8 +811,6 @@ async def test_concurrent_initialize_session_managers(): """Test that concurrent calls to initialize_session_managers don't cause race conditions.""" try: from litellm.proxy._experimental.mcp_server.server import ( - _INITIALIZATION_LOCK, - _SESSION_MANAGERS_INITIALIZED, initialize_session_managers, ) except ImportError: @@ -1426,7 +1528,6 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): ) from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, - LiteLLM_TeamTable, UserAPIKeyAuth, ) except ImportError: diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 78e07c2967..33130d50cc 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -35,7 +35,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, """ mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2]) - # Create mock queues - only 3 of 7 have data + # Create mock queues - only 3 of 6 have data spend_update_queue = AsyncMock() spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( return_value={"key_list_transactions": {"key1": 1.0}} @@ -67,11 +67,6 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, return_value={} ) - daily_tag_queue = AsyncMock() - daily_tag_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( - return_value={} - ) - await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=spend_update_queue, daily_spend_update_queue=daily_spend_queue, @@ -79,7 +74,6 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, daily_org_spend_update_queue=daily_org_queue, daily_end_user_spend_update_queue=daily_end_user_queue, daily_agent_spend_update_queue=daily_agent_queue, - daily_tag_spend_update_queue=daily_tag_queue, ) # Should be called exactly once (pipeline) @@ -117,7 +111,6 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( daily_org_spend_update_queue=empty_daily_queue, daily_end_user_spend_update_queue=empty_daily_queue, daily_agent_spend_update_queue=empty_daily_queue, - daily_tag_spend_update_queue=empty_daily_queue, ) mock_redis_cache.async_rpush_pipeline.assert_not_called() @@ -131,7 +124,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( Verify get_all_transactions_from_redis_buffer_pipeline correctly parses and aggregates results from async_lpop_pipeline. """ - # Simulate pipeline results: slot 0 = spend updates, slots 1-6 = daily categories + # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories db_spend_json = json.dumps( { "key_list_transactions": {"key1": 1.0, "key2": 2.0}, @@ -154,14 +147,13 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( None, # slot 3: daily org (empty) None, # slot 4: daily end-user (empty) None, # slot 5: daily agent (empty) - None, # slot 6: daily tag (empty) ] ) result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() - assert len(result) == 7 - db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent, daily_tag = result + assert len(result) == 6 + db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent = result # Verify db spend was parsed correctly assert db_spend is not None @@ -181,7 +173,6 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( assert daily_org is None assert daily_end_user is None assert daily_agent is None - assert daily_tag is None # Verify pipeline was called once with correct keys mock_redis_cache.async_lpop_pipeline.assert_called_once() @@ -192,7 +183,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): """When redis_cache is None, should return all Nones""" buffer = RedisUpdateBuffer(redis_cache=None) result = await buffer.get_all_transactions_from_redis_buffer_pipeline() - assert result == (None, None, None, None, None, None, None) + assert result == (None, None, None, None, None, None) def test_validate_redis_transaction_buffer_raises_without_redis():