fix(mcp): roll up MCP tool spend to user counters and usage UI by Sameerlite · Pull Request #31576 · BerriAI/litellm

Sameerlite · 2026-06-29T06:13:51Z

Summary

Fire MCP success logging on /mcp-rest/tools/call so direct REST MCP tool calls write spend_logs and update key/user/team spend counters.
Enrich MCP logging metadata with full user/team/org IDs before cost rollup (REST path, Playground path, and cost callback).
Include per-session MCP spend/count in spend logs UI responses.

Additional MCP spend-log fixes

Streaming MCP completions missing spend logs: MCPStreamingIterator returned on the final chunk without draining the inner CustomStreamWrapper, so end-of-stream success handlers never ran. Drain the inner stream after processing tool calls so list_mcp_tools, call_mcp_tool, and the parent completion all write to LiteLLM_SpendLogs.
list_mcp_tools logging crash: async_success_handler received raw MCPTool objects and failed with Object of type Tool is not JSON serializable. Serialize tools with model_dump(mode="json") before logging.
Missing request_tags on MCP sub-calls: list_mcp_tools and call_mcp_tool spend logs had empty request_tags because tag extraction only checked metadata.tags. Reuse StandardLoggingPayloadSetup._get_request_tags so parent tags (including User-Agent-derived tags from proxy_server_request) propagate to MCP sub-call logs.

Test plan

Streaming chat completion with MCP tools writes spend logs for LLM call, list_mcp_tools, and call_mcp_tool
list_mcp_tools logs succeed without JSON serialization errors
MCP sub-call spend logs inherit parent request_tags

Note

Medium Risk
Touches spend logging, cost callbacks, and a raw SQL enrichment query scoped by api_key; failures are mostly non-blocking, but incorrect rollup metadata could mis-attribute MCP spend.

Overview
This PR tightens MCP spend logging and rollups so direct REST calls, gateway sub-calls, and streaming completions all land in LiteLLM_SpendLogs and user/team counters.

/mcp-rest/tools/call now runs shared _fire_mcp_success_logging after tool execution (including the virtual tool-search path), wrapped in _safe_fire_mcp_success_logging so logging failures only warn and never fail the request. call_mcp_tool uses the same helper instead of inlined logging.

Cost tracking enriches metadata with user/team/org when only user_api_key is present (typical MCP REST), via a conditional lookup in _PROXY_track_cost_callback and _write_spend_metadata_to_kwargs. MCP call_mcp_tool logging uses LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata for fuller metadata.

Parent request tags propagate to list_mcp_tools and call_mcp_tool through _get_parent_request_tags (same logic as standard logging, including User-Agent). list_mcp_tools success logging serializes tools with model_dump(mode="json") to avoid JSON errors.

Streaming MCP chat completions drain the inner CustomStreamWrapper after the final chunk so end-of-stream spend handlers run; drain errors are swallowed so the final chunk still reaches the client.

The spend logs UI adds per-session mcp_tool_call_count and mcp_tool_call_spend, scoped by api_key to avoid cross-tenant leakage on colliding session_ids.

^{Reviewed by Cursor Bugbot for commit 2622db2. Bugbot is set up for automated code reviews on this repo. Configure here.}

Direct REST MCP tool calls now fire success logging so spend_logs and user/team rollups include configured mcp_server_cost_info charges. Co-authored-by: Cursor <cursoragent@cursor.com>

greptile-apps · 2026-06-29T06:18:20Z

Greptile Summary

This PR wires MCP tool-call spend logging end-to-end across three entry paths (direct REST, Playground, streaming chat) and adds per-session MCP cost aggregates to the spend-logs UI. The changes are self-contained to the MCP subsystem and its cost callback.

REST MCP logging: /mcp-rest/tools/call now calls _fire_mcp_success_logging after tool execution via the new _safe_fire_mcp_success_logging helper; logging failures are caught and warned rather than propagating to callers.
Metadata enrichment: _PROXY_track_cost_callback conditionally fetches user/team/org IDs from the key object when user_api_key_user_id is absent, enabling rollup for MCP calls that only carry a hashed key.
Streaming drain: The inner CustomStreamWrapper is now drained after the final chunk so its end-of-stream spend handlers fire; drain errors are swallowed and logged so the already-assembled final chunk is never dropped.
UI aggregation: A new parameterized query_raw in _build_ui_spend_logs_response joins MCP sub-call spend to spend-log rows, scoped by api_key to prevent cross-tenant leakage on colliding session_ids.

Confidence Score: 5/5

Safe to merge. All changes are isolated to MCP spend-logging paths; failures are wrapped in try/except or in _safe_fire_mcp_success_logging so logging errors cannot break the tool-call response. The new SQL query is parameterized and scoped by api_key to prevent cross-tenant leakage.

The core logic — firing success logging on the REST path, enriching metadata before cost rollup, draining the inner stream, and aggregating MCP spend in the UI — is correct and well-tested. Error handling is consistently defensive: logging failures warn and continue, the drain swallows non-StopAsyncIteration exceptions, and the UI enrichment query degrades silently on Prisma errors. No rollup mis-attribution or double-counting path was found.

No files require special attention. litellm/proxy/hooks/proxy_track_cost_callback.py has a minor write-back gap in _write_spend_metadata_to_kwargs when neither metadata bucket exists, but cost tracking itself is unaffected.

Important Files Changed

Filename	Overview
litellm/proxy/_experimental/mcp_server/rest_endpoints.py	Adds `_safe_fire_mcp_success_logging` helper and fires it after both the virtual-tool and standard execute paths; `CancelledError` is correctly re-raised while other errors warn-and-continue.
litellm/proxy/_experimental/mcp_server/server.py	Extracts `_fire_mcp_success_logging` into a standalone helper; `call_mcp_tool` now calls it via the helper instead of inlining; `list_mcp_tools` serialises tools with `model_dump(mode="json")` before logging and accepts `request_tags`.
litellm/proxy/hooks/proxy_track_cost_callback.py	Adds conditional `_enrich_failure_metadata_with_key_info` call in the success callback when `user_api_key_user_id` is absent, enabling user/team rollup for MCP REST calls; adds `_write_spend_metadata_to_kwargs` to propagate enriched fields back to kwargs.
litellm/proxy/spend_tracking/spend_management_endpoints.py	New `query_raw` in `_build_ui_spend_logs_response` aggregates MCP sub-call counts and spend per session; scoped by `api_key` collected from the authorized page rows; errors are swallowed at debug level.
litellm/responses/mcp/chat_completions_handler.py	Adds `_drain_inner_stream` to exhaust the `CustomStreamWrapper` after the final chunk so spend handlers fire; drain exceptions are caught and logged, preserving the final chunk; `request_tags` are threaded through the streaming iterator.
litellm/responses/mcp/litellm_proxy_mcp_handler.py	Adds `_get_parent_request_tags` static method reusing `StandardLoggingPayloadSetup._get_request_tags`; upgrades `_execute_tool_calls` to use `LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata` for fuller metadata; threads `request_tags` through all call sites.
litellm/responses/mcp/mcp_streaming_iterator.py	Minimal change: propagates `request_tags` to `_execute_tool_calls` in `MCPEnhancedStreamingIterator`.
litellm/responses/main.py	Threads `request_tags` from `_get_parent_request_tags` to all three MCP tool-processing call sites in `aresponses_api_with_mcp`.

_{Reviews (8): Last reviewed commit: "fix: propagate MCP logging cancellation" | Re-trigger Greptile}

codecov · 2026-06-29T06:19:24Z

Codecov Report

❌ Patch coverage is 90.66667% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/hooks/proxy_track_cost_callback.py	85.71%	2 Missing ⚠️
...proxy/spend_tracking/spend_management_endpoints.py	84.61%	2 Missing ⚠️
...m/proxy/_experimental/mcp_server/rest_endpoints.py	93.33%	1 Missing ⚠️
litellm/proxy/_experimental/mcp_server/server.py	83.33%	1 Missing ⚠️
litellm/responses/mcp/litellm_proxy_mcp_handler.py	92.85%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

…mport order - Only call _enrich_failure_metadata_with_key_info when user_api_key_user_id is absent, avoiding a cache/DB lookup on every normal LLM request. - Move LiteLLMProxyRequestSetup import to correct alphabetical position (I001). Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-06-29T07:46:01Z

@greptileai

veria-ai · 2026-06-29T07:55:50Z

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

…t disclosure Add api_key = ANY($2) to the MCP session aggregate query so it is bounded by the same ownership already applied to the main page query. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-06-29T09:36:20Z

@greptileai

Replace the 8 new UP006 violations introduced by the mcp-tags changes: - Optional[List[str]] → Optional[list[str]] for request_tags params - List[str] return type → list[str] in _get_parent_request_tags - Dict[str, Dict[...]] → dict[str, dict[...]] for mcp_spend_map annotation Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-06-30T04:42:18Z

@greptileai

mateo-berri · 2026-07-02T06:32:24Z

@greptileai

Generated by Claude Code

…w MCP spend enrichment except to PrismaError

mateo-berri · 2026-07-02T06:58:16Z

@greptileai

Generated by Claude Code

mateo-berri · 2026-07-02T07:24:03Z

@greptileai

Generated by Claude Code

mateo-berri · 2026-07-02T07:35:51Z

bugbot run

Generated by Claude Code

cursor

Cursor Bugbot has reviewed your changes using high effort and found 4 potential issues.

Autofix Details

Bugbot Autofix prepared fixes for all 4 issues found in the latest run.

✅ Fixed: REST logging errors fail requests
- Wrapped REST MCP success logging in a non-throwing helper so completed tool calls still return if logging fails.
✅ Fixed: Stream drain skipped on exhaustion
- Added inner stream draining on the StopAsyncIteration path after processing the final collected MCP chat chunk.
✅ Fixed: Parent tags use wrong params
- Changed parent tag extraction to read nested litellm_params and proxy_server_request while preserving the top-level fallback.
✅ Fixed: Virtual REST path skips logging
- The virtual REST mcp_tool_call branch now fires the same isolated success logging path before returning the tool result.

_{You can send follow-ups to the cloud agent here.}

CLAassistant · 2026-07-02T07:56:09Z

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
2 out of 3 committers have signed the CLA.

✅ Sameerlite
✅ mateo-berri
❌ cursoragent
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

Sameerlite · 2026-07-02T08:25:15Z

bugbot run

cursor

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

Bugbot Autofix prepared a fix for the issue found in the latest run.

✅ Fixed: Safe logging swallows cancellation
- Cancellation from MCP success logging is now re-raised before ordinary logging failures are swallowed, with a regression test covering the propagation.

_{You can send follow-ups to the cloud agent here.}

^{Reviewed by Cursor Bugbot for commit 2622db2. Configure here.}

mateo-berri · 2026-07-02T15:04:57Z

@greptileai

Generated by Claude Code

mateo-berri

LGTM; thanks!

Generated by Claude Code

…AI#31576) * fix(mcp): roll up MCP tool spend to user counters and usage UI Direct REST MCP tool calls now fire success logging so spend_logs and user/team rollups include configured mcp_server_cost_info charges. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(mcp): gate key-info enrichment to requests missing user_id; fix import order - Only call _enrich_failure_metadata_with_key_info when user_api_key_user_id is absent, avoiding a cache/DB lookup on every normal LLM request. - Move LiteLLMProxyRequestSetup import to correct alphabetical position (I001). Co-authored-by: Cursor <cursoragent@cursor.com> * fix(mcp): scope MCP spend aggregate by api_key to prevent cross-tenant disclosure Add api_key = ANY($2) to the MCP session aggregate query so it is bounded by the same ownership already applied to the main page query. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix spend logs for call and list mcp tools * Add tags in mcp logging * Fix ruff * fix(lint): replace List/Dict with list/dict in new annotations (UP006) Replace the 8 new UP006 violations introduced by the mcp-tags changes: - Optional[List[str]] → Optional[list[str]] for request_tags params - List[str] return type → list[str] in _get_parent_request_tags - Dict[str, Dict[...]] → dict[str, dict[...]] for mcp_spend_map annotation Co-authored-by: Cursor <cursoragent@cursor.com> * fix(lint): keep call_tool_rest_api within complexity budget and narrow MCP spend enrichment except to PrismaError * fix(mcp): keep final streaming chunk when draining inner stream fails * fix: handle MCP logging edge cases * fix: propagate MCP logging cancellation --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

## What - Add `orjson>=3.11.6,<4.0` as a direct SDK runtime dependency for tool-enabled rollout calls across local and remote backends. - Regenerate `uv.lock`, resolving `orjson==3.11.9`. ## Why Tool-enabled rollout calls can fail on any backend that installs the base SDK dependency set without separately installing `orjson`. This affected both `LocalBackend` and Harbor/Daytona sandbox execution, failing with `litellm.APIConnectionError: OpenAIException - No module named 'orjson'`. LiteLLM imports `LiteLLM_Proxy_MCP_Handler` whenever a completion contains any `tools` ([v1.92.0 source](https://github.com/BerriAI/litellm/blob/v1.92.0/litellm/main.py#L4865-L4876)). [LiteLLM PR #31576](BerriAI/litellm#31576) added a top-level `LiteLLMProxyRequestSetup` import to that handler ([commit](BerriAI/litellm@fabe5c2)), which imports `http_parsing_utils` and therefore `orjson` at module load time. LiteLLM still declares `orjson` only under its `proxy` optional extra ([pyproject.toml](https://github.com/BerriAI/litellm/blob/v1.92.0/pyproject.toml#L42-L54)), so a base LiteLLM installation can reach this runtime path without having `orjson` installed. Declaring `orjson` directly with LiteLLM's existing version bound makes the SDK dependency set complete for tool-enabled rollout calls across local and remote backends. Adding `orjson` to the rollout image also unblocked the failed remote evaluation and artifact collection, confirming that the dependency gap—not artifact persistence—caused that incident. ## How to Test - `uv lock --check` - `uv run --locked --extra dev ruff check .` - `uv run --locked --extra dev ruff format --check .` - `uv run --locked --extra dev pyright osmosis_ai/` - `uv run --locked --extra dev pytest -q` - `uv build --wheel --out-dir /tmp/osmosis-sdk-orjson-dist` - `unzip -p /tmp/osmosis-sdk-orjson-dist/*.whl '*/METADATA' | rg '^Requires-Dist: orjson'` ## Checklist - [x] PR title follows `[module] type: description` format - [x] Appropriate labels added (`bug`, `rollout`, `dependencies`) - [x] `ruff check .` and `ruff format --check .` pass - [x] `pyright osmosis_ai/` passes - [x] `pytest` passes (`1607 passed`) - [x] Public API changes are documented (no public API changes) - [x] No secrets or credentials included

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

cca1056

Direct REST MCP tool calls now fire success logging so spend_logs and user/team rollups include configured mcp_server_cost_info charges. Co-authored-by: Cursor <cursoragent@cursor.com>

greptile-apps Bot reviewed Jun 29, 2026

View reviewed changes

Comment thread litellm/proxy/hooks/proxy_track_cost_callback.py Outdated

veria-ai Bot reviewed Jun 29, 2026

View reviewed changes

Comment thread litellm/proxy/spend_tracking/spend_management_endpoints.py

fix(mcp): scope MCP spend aggregate by api_key to prevent cross-tenan…

e08ef3b

…t disclosure Add api_key = ANY($2) to the MCP session aggregate query so it is bounded by the same ownership already applied to the main page query. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite and others added 4 commits June 30, 2026 09:47

Fix spend logs for call and list mcp tools

74a44af

Add tags in mcp logging

5bfc365

Fix ruff

87e68ef

merge litellm_internal_staging into litellm_mcp-cost-tracking-rollup

956ed23

fix(lint): keep call_tool_rest_api within complexity budget and narro…

ac7d5d0

…w MCP spend enrichment except to PrismaError

fix(mcp): keep final streaming chunk when draining inner stream fails

0e0ba0f

cursor Bot reviewed Jul 2, 2026

View reviewed changes

Comment thread litellm/proxy/_experimental/mcp_server/rest_endpoints.py Outdated

Comment thread litellm/responses/mcp/chat_completions_handler.py

Comment thread litellm/responses/mcp/litellm_proxy_mcp_handler.py

Comment thread litellm/proxy/_experimental/mcp_server/rest_endpoints.py

fix: handle MCP logging edge cases

2622db2

cursor Bot reviewed Jul 2, 2026

View reviewed changes

Comment thread litellm/proxy/_experimental/mcp_server/rest_endpoints.py

fix: propagate MCP logging cancellation

b167e62

mateo-berri approved these changes Jul 2, 2026

View reviewed changes

mateo-berri merged commit fabe5c2 into litellm_internal_staging Jul 2, 2026
121 of 122 checks passed

mateo-berri deleted the litellm_mcp-cost-tracking-rollup branch July 2, 2026 15:17

This was referenced Jul 7, 2026

fix(mcp): defer proxy import so completion(tools=...) works without proxy extras #32339

Merged

ci(codspeed): re-enable benchmarks on litellm_internal_staging #32340

Merged

JoyboyBrian mentioned this pull request Jul 13, 2026

[rollout] fix: add missing orjson runtime dependency Osmosis-AI/osmosis-sdk-python#255

Merged

7 tasks

Uh oh!

Uh oh!

Conversation

Sameerlite commented Jun 29, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Additional MCP spend-log fixes

Test plan

Uh oh!

greptile-apps Bot commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Uh oh!

Uh oh!

codecov Bot commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Sameerlite commented Jun 29, 2026

Uh oh!

Uh oh!

veria-ai Bot commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR overview

Security review

Uh oh!

Sameerlite commented Jun 29, 2026

Uh oh!

Sameerlite commented Jun 30, 2026

Uh oh!

mateo-berri commented Jul 2, 2026

Uh oh!

mateo-berri commented Jul 2, 2026

Uh oh!

mateo-berri commented Jul 2, 2026

Uh oh!

mateo-berri commented Jul 2, 2026

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

CLAassistant commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Sameerlite commented Jul 2, 2026

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mateo-berri commented Jul 2, 2026

Uh oh!

mateo-berri left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Sameerlite commented Jun 29, 2026 •

edited by cursor Bot

Loading

greptile-apps Bot commented Jun 29, 2026 •

edited

Loading

codecov Bot commented Jun 29, 2026 •

edited

Loading

veria-ai Bot commented Jun 29, 2026 •

edited

Loading

cursor Bot left a comment •

edited

Loading

CLAassistant commented Jul 2, 2026 •

edited

Loading

cursor Bot left a comment •

edited

Loading