fix(mcp): roll up MCP tool spend to user counters and usage UI

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -77,6 +77,7 @@ def _connection_error_message(exc: BaseException) -> str:
         ListMCPToolsRestAPIResponseObject,
         MCPInfo,
         MCPServer,
+        _fire_mcp_success_logging,
         _tool_name_matches,
         execute_mcp_tool,
         filter_tools_by_allowed_tools,
@@ -796,11 +797,12 @@ async def call_tool_rest_api(
                 user_oauth_extra_headers = await _get_user_oauth_extra_headers(target_server, user_api_key_dict)
 
             # Call execute_mcp_tool directly (permission checks already done)
+            _tool_start_time = datetime.now()
             result = await execute_mcp_tool(
                 name=tool_name,
                 arguments=tool_arguments,
                 allowed_mcp_servers=allowed_mcp_servers,
-                start_time=datetime.now(),
+                start_time=_tool_start_time,
                 user_api_key_auth=data.get("user_api_key_auth"),
                 mcp_auth_header=data.get("mcp_auth_header"),
                 mcp_server_auth_headers=data.get("mcp_server_auth_headers"),
@@ -809,6 +811,8 @@ async def call_tool_rest_api(
                 litellm_logging_obj=data.get("litellm_logging_obj"),
                 requested_server_id=canonical_server_id,
             )
+            if logging_obj is not None:
+                await _fire_mcp_success_logging(logging_obj, result, _tool_start_time, datetime.now())
             return result
         except MCPMissingUserEnvVarsError as e:
             verbose_logger.info(

litellm/proxy/_experimental/mcp_server/server.py

@@ -2526,6 +2526,22 @@ async def execute_mcp_tool(
 
         return response
 
+    async def _fire_mcp_success_logging(
+        logging_obj: LiteLLMLoggingObj,
+        result: Any,
+        start_time: datetime,
+        end_time: datetime,
+    ) -> None:
+        logging_obj.post_call(original_response=result)
+        await logging_obj.async_post_mcp_tool_call_hook(
+            kwargs=logging_obj.model_call_details,
+            response_obj=result,
+            start_time=start_time,
+            end_time=end_time,
+        )
+        logging_obj.call_type = CallTypes.call_mcp_tool.value
+        await logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time)
+
     @client
     async def call_mcp_tool(
         name: str,
@@ -2597,16 +2613,7 @@ async def call_mcp_tool(
             raise
 
         if litellm_logging_obj:
-            litellm_logging_obj.post_call(original_response=response)
-            end_time = datetime.now()
-            await litellm_logging_obj.async_post_mcp_tool_call_hook(
-                kwargs=litellm_logging_obj.model_call_details,
-                response_obj=response,
-                start_time=start_time,
-                end_time=end_time,
-            )
-            litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value
-            await litellm_logging_obj.async_success_handler(result=response, start_time=start_time, end_time=end_time)
+            await _fire_mcp_success_logging(litellm_logging_obj, response, start_time, datetime.now())
         return response
 
     async def mcp_get_prompt(

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -190,6 +190,11 @@ async def _PROXY_track_cost_callback(
             litellm_params = kwargs.get("litellm_params", {}) or {}
             end_user_id = get_end_user_id_for_cost_tracking(litellm_params)
             metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
+            # Only fetch key details when user_id wasn't already populated (e.g. direct MCP REST calls).
+            # Avoids a cache/DB lookup on every normal LLM request.
+            if metadata.get("user_api_key") and not metadata.get("user_api_key_user_id"):
+                metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata)
+                _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
             budget_reservation = _get_budget_reservation_from_metadata(metadata=metadata)
             user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None))
             team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None))
@@ -388,6 +393,20 @@ def _should_track_errors_in_db():
         return
 
 
+def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None:
+    patch = {k: v for k, v in metadata.items() if (k.startswith("user_api_key") or k == "tags") and v is not None}
+    if not patch:
+        return
+
+    litellm_params = kwargs.setdefault("litellm_params", {})
+    for bucket_name in ("litellm_metadata", "metadata"):
+        bucket = litellm_params.get(bucket_name)
+        if isinstance(bucket, dict):
+            for key, value in patch.items():
+                if bucket.get(key) is None:
+                    bucket[key] = value
+
+
 def _should_track_cost_callback(
     user_api_key: Optional[str],
     user_id: Optional[str],

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -3444,12 +3444,45 @@ async def _build_ui_spend_logs_response(
             )
             count_map = {r["session_id"]: r["_count"]["session_id"] for r in counts if r.get("session_id")}
 
+    mcp_spend_map: Dict[str, Dict[str, Union[int, float]]] = {}
+    if enrich_session_counts and session_ids:
+        try:
+            rows = await prisma_client.db.query_raw(
+                """
+                SELECT session_id,
+                       COUNT(*)::int AS mcp_tool_call_count,
+                       COALESCE(SUM(spend), 0)::double precision AS mcp_tool_call_spend
+                FROM "LiteLLM_SpendLogs"
+                WHERE session_id = ANY($1::text[])
+                  AND call_type IN ('call_mcp_tool', 'list_mcp_tools')
+                GROUP BY session_id
+                """,
+                session_ids,
+            )
+            mcp_spend_map = {
+                row["session_id"]: {
+                    "mcp_tool_call_count": int(row.get("mcp_tool_call_count") or 0),
+                    "mcp_tool_call_spend": float(row.get("mcp_tool_call_spend") or 0.0),
+                }
+                for row in rows
+                if row.get("session_id")
+            }
+        except Exception:
+            verbose_proxy_logger.debug(
+                "Failed to enrich MCP session spend aggregates for spend logs UI",
+                exc_info=True,
+            )
+
     if enrich_session_counts:
         enriched: List[dict] = []
         for row in data:
             row_dict = dict(row) if isinstance(row, dict) else row.model_dump()
             sid = row_dict.get("session_id")
             row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1
+            mcp_stats = mcp_spend_map.get(sid) if sid else None
+            if mcp_stats:
+                row_dict["mcp_tool_call_count"] = mcp_stats["mcp_tool_call_count"]
+                row_dict["mcp_tool_call_spend"] = mcp_stats["mcp_tool_call_spend"]
             enriched.append(row_dict)
         response_data: list = enriched
     else:

litellm/responses/mcp/litellm_proxy_mcp_handler.py

@@ -17,6 +17,7 @@
 from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy._experimental.mcp_server.utils import split_server_prefix_from_name
+from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
 from litellm.responses.main import aresponses
 from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
 from litellm.types.llms.openai import ResponsesAPIResponse
@@ -672,17 +673,17 @@ async def _execute_tool_calls(
                 }
                 if litellm_trace_id:
                     logging_request_data["litellm_trace_id"] = litellm_trace_id
-                user_identifier = None
                 if user_api_key_auth is not None:
-                    user_api_key = getattr(user_api_key_auth, "api_key", None)
-                    if user_api_key:
-                        logging_request_data["metadata"]["user_api_key"] = user_api_key
-
+                    LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
+                        data=logging_request_data,
+                        user_api_key_dict=user_api_key_auth,
+                        _metadata_variable_name="metadata",
+                    )
                     user_identifier = getattr(user_api_key_auth, "end_user_id", None) or getattr(
                         user_api_key_auth, "user_id", None
                     )
-                if user_identifier:
-                    logging_request_data["user"] = user_identifier
+                    if user_identifier:
+                        logging_request_data["user"] = user_identifier
 
                 litellm_logging_obj: Optional[LiteLLMLoggingObj] = None
                 try:

tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py

@@ -1,6 +1,6 @@
 import json
 from typing import Any, Dict, Optional
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import httpx
 import pytest
@@ -330,7 +330,8 @@ async def ok_operation(client):
     @pytest.mark.asyncio
     async def test_m2m_does_not_build_presented_store(self, monkeypatch):
         """M2M (client_credentials): to_server_spec returns None, so no presented provider is built;
-        the auto-fetch path is unchanged (no cred_provider, the incoming header dropped as before)."""
+        the auto-fetch path is unchanged (no cred_provider, the incoming header dropped as before).
+        """
         captured: dict = {}
 
         def fake_build_stdio_env(server, raw_headers):
@@ -377,7 +378,8 @@ async def ok_operation(client):
     @pytest.mark.asyncio
     async def test_token_exchange_does_not_build_presented_store(self, monkeypatch):
         """OBO / token-exchange (auth_type oauth2_token_exchange, not oauth2): excluded by the
-        auth_type == oauth2 guard, so no presented provider is built and the v1 exchange path runs."""
+        auth_type == oauth2 guard, so no presented provider is built and the v1 exchange path runs.
+        """
         captured: dict = {}
 
         def fake_build_stdio_env(server, raw_headers):
@@ -1196,6 +1198,13 @@ async def fake_execute_mcp_tool(**kwargs):
             fake_execute_mcp_tool,
             raising=False,
         )
+        fire_logging = AsyncMock()
+        monkeypatch.setattr(
+            rest_endpoints,
+            "_fire_mcp_success_logging",
+            fire_logging,
+            raising=False,
+        )
 
         request_payload = {
             "server_id": "server-1",
@@ -1217,6 +1226,7 @@ async def fake_execute_mcp_tool(**kwargs):
         assert captured["name"] == "demo-tool"
         assert captured["arguments"] == {"foo": "bar"}
         assert captured["allowed_mcp_servers"] == [stub_server]
+        fire_logging.assert_awaited_once()
 
 
 class TestGetToolsForSingleServer:
@@ -1809,9 +1819,9 @@ async def fake_load_spec(spec_path):  # noqa: ANN001
         names = [t["name"] for t in result["tools"]]
         anthropic_re = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")
         for name in names:
-            assert anthropic_re.match(name), (
-                f"preview tool name {name!r} violates ^[a-zA-Z0-9_-]+$"
-            )
+            assert anthropic_re.match(
+                name
+            ), f"preview tool name {name!r} violates ^[a-zA-Z0-9_-]+$"
         assert "actions_download-job-logs-for-workflow-run" in names
         assert "pulls_list-files" in names
 
@@ -1868,7 +1878,9 @@ async def fake_load_spec(spec_path):  # noqa: ANN001
 
         registered_summary_to_name: dict = {}
 
-        def fake_create_tool_function(path, method, operation, base_url):  # noqa: ANN001
+        def fake_create_tool_function(
+            path, method, operation, base_url
+        ):  # noqa: ANN001
             def _f():
                 return None
 
@@ -1881,7 +1893,9 @@ def _f():
         )
 
         class _StubRegistry:
-            def register_tool(self, name, description, input_schema, handler):  # noqa: ANN001
+            def register_tool(
+                self, name, description, input_schema, handler
+            ):  # noqa: ANN001
                 registered_summary_to_name[description] = name
 
         monkeypatch.setattr(

tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py

@@ -1104,3 +1104,76 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend():
 
         mock_update_database.assert_called_once()
         assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05
+
+
+@pytest.mark.asyncio
+async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata():
+    """MCP tool calls may only carry user_api_key; user/team rollups still need user_id."""
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    logger = _ProxyDBLogger()
+    key_obj = UserAPIKeyAuth(
+        api_key="hashed-key",
+        user_id="mcp-user@example.com",
+        team_id="team-123",
+        org_id="org-456",
+        key_alias="mcp-key",
+    )
+
+    kwargs = {
+        "call_type": "call_mcp_tool",
+        "model": "MCP: echo",
+        "litellm_params": {
+            "metadata": {
+                "user_api_key": "hashed-key",
+            }
+        },
+        "standard_logging_object": {
+            "response_cost": 10.0,
+            "request_tags": [],
+            "metadata": {},
+        },
+    }
+
+    with (
+        patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+            new_callable=AsyncMock,
+            return_value=key_obj,
+        ),
+        patch(
+            "litellm.proxy.proxy_server.increment_spend_counters",
+            new_callable=AsyncMock,
+        ) as mock_increment,
+        patch(
+            "litellm.proxy.proxy_server.update_cache",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.proxy_server.proxy_logging_obj",
+        ) as mock_proxy_logging,
+    ):
+        mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
+        mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
+
+        await logger._PROXY_track_cost_callback(
+            kwargs=kwargs,
+            completion_response={"id": "mcp-call-1"},
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+        )
+
+        mock_increment.assert_awaited_once()
+        assert mock_increment.call_args.kwargs["user_id"] == "mcp-user@example.com"
+        assert mock_increment.call_args.kwargs["team_id"] == "team-123"
+        assert mock_increment.call_args.kwargs["org_id"] == "org-456"
+
+        update_kwargs = (
+            mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs
+        )
+        assert update_kwargs["user_id"] == "mcp-user@example.com"
+        assert update_kwargs["team_id"] == "team-123"
+        assert (
+            kwargs["litellm_params"]["metadata"]["user_api_key_user_id"]
+            == "mcp-user@example.com"
+        )

tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py

@@ -1051,7 +1051,10 @@ async def mock_query_raw(sql_query, *params):
         page_size = params[-2] if len(params) >= 2 else 50
         skip = params[-1] if len(params) >= 1 else 0
         return [
-            {**{k: v for k, v in row.items() if k != "_ttft_ms"}, "total_count": len(base_logs)}
+            {
+                **{k: v for k, v in row.items() if k != "_ttft_ms"},
+                "total_count": len(base_logs),
+            }
             for row in sorted_logs[skip : skip + page_size]
         ]
 
@@ -2929,6 +2932,15 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
             {"session_id": session_id, "_count": {"session_id": 2}},
         ]
     )
+    mock_prisma.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "session_id": session_id,
+                "mcp_tool_call_count": 1,
+                "mcp_tool_call_spend": 10.0,
+            }
+        ]
+    )
 
     result = await _build_ui_spend_logs_response(
         prisma_client=mock_prisma,
@@ -2946,6 +2958,10 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
     # Rows with the shared session_id should have session_total_count=2
     assert rows[0]["session_total_count"] == 2
     assert rows[1]["session_total_count"] == 2
+    assert rows[0]["mcp_tool_call_count"] == 1
+    assert rows[0]["mcp_tool_call_spend"] == 10.0
+    assert rows[1]["mcp_tool_call_count"] == 1
+    assert rows[1]["mcp_tool_call_spend"] == 10.0
 
     # Row without a session_id defaults to 1
     assert rows[2]["session_total_count"] == 1
@@ -4104,7 +4120,9 @@ async def test_cold_storage_handler_returns_none_when_no_logger_configured(monke
 
 
 @pytest.mark.asyncio
-async def test_cold_storage_handler_resolves_configured_logger_from_registry(monkeypatch):
+async def test_cold_storage_handler_resolves_configured_logger_from_registry(
+    monkeypatch,
+):
     from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
 
     logger = _FakeColdStorageLogger({"messages": "from-registry"})