From cc3ffbe89c707a768ccf49e7888e31b3446a7c67 Mon Sep 17 00:00:00 2001
From: sathvikvittal <sathvik.vittal@gmail.com>
Date: Mon, 10 Aug 2026 11:49:14 +0530
Subject: [PATCH] fix(proxy): separate daily spend by MCP tool

---
 .../proxy/db/test_db_spend_update_writer.py   | 51 ++++++++++++++++---
 1 file changed, 45 insertions(+), 6 deletions(-)

diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py
index 03e28a879fc..1bc08cc8118 100644
--- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py
+++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py
@@ -12,7 +12,7 @@ sys.path.insert(
 
 from collections.abc import Callable
 from datetime import datetime, timezone
-from unittest.mock import AsyncMock, MagicMock, call, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from redis.exceptions import DataError
@@ -375,6 +375,47 @@ async def test_daily_user_transaction_key_includes_model_group():
     ]
 
 
+@pytest.mark.asyncio
+async def test_daily_user_transaction_key_includes_mcp_tool_name():
+    writer = DBSpendUpdateWriter()
+    writer.daily_spend_update_queue = MagicMock()
+    writer.daily_spend_update_queue.add_update = AsyncMock()
+    mock_prisma = MagicMock()
+
+    base_payload = {
+        "request_id": "req-mcp-tool-key",
+        "user": "test-user",
+        "startTime": "2024-01-01T12:00:00",
+        "api_key": "test-key",
+        "model": "openai/gpt-4o-mini",
+        "model_group": "a",
+        "custom_llm_provider": "openai",
+        "prompt_tokens": 10,
+        "completion_tokens": 5,
+        "spend": 0.2,
+        "metadata": '{"usage_object": {}}',
+    }
+
+    await writer.add_spend_log_transaction_to_daily_user_transaction(
+        payload={**base_payload, "mcp_namespaced_tool_name": "server_a.tool"},
+        prisma_client=mock_prisma,
+    )
+    await writer.add_spend_log_transaction_to_daily_user_transaction(
+        payload={**base_payload, "mcp_namespaced_tool_name": "server_b.tool"},
+        prisma_client=mock_prisma,
+    )
+
+    queued_keys = [
+        next(iter(call_kwargs.kwargs["update"].keys()))
+        for call_kwargs in writer.daily_spend_update_queue.add_update.call_args_list
+    ]
+
+    assert queued_keys == [
+        '["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_a.tool",""]',
+        '["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_b.tool",""]',
+    ]
+
+
 def test_daily_transaction_key_is_not_delimiter_ambiguous():
     assert DBSpendUpdateWriter._daily_transaction_key(
         "user", "model_a", "group"
@@ -677,7 +718,7 @@ async def test_update_tag_db_with_valid_tags():
     """
     Test that _update_tag_db correctly processes valid tags and adds them to the spend update queue.
     """
-    from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem
+    from litellm.proxy._types import Litellm_EntityType
 
     writer = DBSpendUpdateWriter()
     mock_prisma = MagicMock()
@@ -1019,8 +1060,6 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
         "metadata": '{"usage_object": {}}',
     }
 
-    # Mock the add_update method to capture what's being added
-    original_add_update = writer.daily_tag_spend_update_queue.add_update
     writer.daily_tag_spend_update_queue.add_update = AsyncMock()
 
     await writer.add_spend_log_transaction_to_daily_tag_transaction(
@@ -1032,8 +1071,8 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
     assert writer.daily_tag_spend_update_queue.add_update.call_count == 2
 
     # Check that request_id is included in both transactions
-    for call in writer.daily_tag_spend_update_queue.add_update.call_args_list:
-        transaction_dict = call[1]["update"]
+    for update_call in writer.daily_tag_spend_update_queue.add_update.call_args_list:
+        transaction_dict = update_call[1]["update"]
         # Each transaction should have one key with the format tag_date_api_key_model_provider
         for key, transaction in transaction_dict.items():
             assert (